chore: bump to 1.3.10

fix: memmap mode=r, tool parser json.loads, greedy decode
refactor: move sample-id indexing from dataset to store
2026-07-19 16:40:27 +08:00 · 2026-07-19 16:38:28 +08:00 · 2026-07-19 16:02:50 +08:00 · 2026-07-19 14:49:58 +08:00 · 2026-07-19 12:45:58 +08:00 · 2026-07-19 12:33:35 +08:00
152 changed files with 22654 additions and 5261 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/**
@@ -21,3 +29,8 @@
 !/LICENSE
 !/pyproject.toml
 !/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:
+```bash
-  ```bash
+ruff format .
-  ruff format .
+```
  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
-  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>
@@ -20,7 +20,7 @@
  <a href="assets/docs/README-zh-CN.md">中文</a> •
  <a href="https://github.com/ViperEkura/AstrAI/issues">Issue Tracker</a> •
  <a href="https://github.com/ViperEkura/AstrAI/discussions">Discussions</a> •
-  <a href="https://huggingface.co/ViperEk/">HuggingFace</a>
+  <a href="https://huggingface.co/ViperEkura">HuggingFace</a>
 </div>
 <br>
@@ -28,7 +28,8 @@
 ## 📖 Table of Contents
 - [Features](#features)
- [Quick Start](#quick-start)
+- [Getting Started](#getting-started)
 - [Demo](#demo)
 - [Documentation](#documentation)
 - [Contributing](#contributing)
 - [Community](#community)
@@ -49,45 +50,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/
 ```
-#### Train a Model
+**3. Preprocess data**
 Create `pretrain.json` (preprocessing config for `seq` strategy):
 ```json
 {
    "version": 1,
    "input": {"sections": [{"field": "text", "action": "train"}]},
    "preprocessing": {"max_seq_len": 2048},
    "output": {"storage_format": "bin"}
 }
 ```
 ```bash
-CUDA_VISIBLE_DEVICES=0,1,2,3 python scripts/tools/train.py \
+python scripts/tools/preprocess.py data/*.jsonl -o output/ -c pretrain.json
    --train_type seq \
    --data_root_path /path/to/dataset \
    --param_path /path/to/model \
    --batch_size 4 \
    --accumulation_steps 8 \
    --max_lr 3e-4 \
    --warmup_steps 1000 \
    --n_epoch 1
 ```
-Full reference at [Parameter Guide](assets/docs/params.md).
+**4. Train**
-#### Generate Text
+```bash
 export CUDA_VISIBLE_DEVICES=0,1,2,3
 nohup python scripts/tools/train.py \
    --nprocs=4 \
    --parallel_mode=ddp \
    --train_type=seq \
    --data_root_path=/path/to/dataset \
    --param_path=/path/to/model \
    --batch_per_device=4 \
    --grad_accum_steps=8 \
    --warmup_ratio=0.05 \
    --max_lr=1e-4 \
    --max_grad_norm=1.0 \
    --weight_decay=0.1 \
    --window_size=2048 \
    --ckpt_interval=10000 \
    --ckpt_dir=./checkpoint \
    --random_seed=3407 \
    --label_smoothing=0.05 \
    > out.log 2> err.log &
 ```
 **5. Serve & query**
 ```bash
 # Terminal 1: start server
 python scripts/tools/server.py --param_path ./params --device cuda
 # Terminal 2: query
 curl http://localhost:8000/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{"messages":[{"role":"user","content":"Hello"}],"max_tokens":512}'
 ```
 ### Demo
 Check out the demos in the `scripts/demo/` folder:
 ```bash
 # Download model weights (required before running demos)
 python scripts/demo/download.py                      # model → params/
 # Interactive streaming chat (multi-turn, maintains history)
 python scripts/demo/stream_chat.py
 # Type your message after >>, type !exit to quit
 # Batch generation (5 hardcoded prompts, non-streaming)
 python scripts/demo/generate_batch.py
 # Single-prompt autoregressive streaming
 python scripts/demo/generate_ar.py
 ```
 All generation demos use `temperature=0.8`, `top_p=0.95`, `top_k=50`, `max_tokens=2048` by default and require `params/` to contain model weights (run `download.py` first).
 Watch a video walkthrough on [bilibili](https://www.bilibili.com/video/BV1fuLB6yEj6).
 ---
 See [Documentation](#documentation) for full references beyond the examples above.
 #### Text Generation
 Batch generation from a JSONL file:
 ```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
@@ -101,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
@@ -120,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
@@ -217,7 +241,7 @@ For major changes, please open an issue first to discuss what you would like to
 - **GitHub Issues**: [Issue Tracker](https://github.com/ViperEkura/AstrAI/issues)
 - **Discussions**: [GitHub Discussions](https://github.com/ViperEkura/AstrAI/discussions)
- **HuggingFace**: [Model Hub](https://huggingface.co/ViperEk)
+- **HuggingFace**: [Model Hub](https://huggingface.co/ViperEkura)
 ### License
@@ -15,9 +15,9 @@
 <div align="center">
  <img src="https://img.shields.io/badge/python-3.12+-blue.svg" alt="python">
  <img src="https://img.shields.io/badge/license-GPL--3.0-blue.svg" alt="license">
-  <img src="https://img.shields.io/github/v/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>
@@ -27,14 +27,15 @@
  <a href="#chinese">中文</a> •
  <a href="https://github.com/ViperEkura/AstrAI/issues">问题追踪</a> •
  <a href="https://github.com/ViperEkura/AstrAI/discussions">讨论区</a> •
-  <a href="https://huggingface.co/ViperEk">HuggingFace</a>
+  <a href="https://huggingface.co/ViperEkura">HuggingFace</a>
 </div>
 <br>
 ## 📖 目录
 - [特性](#特性)
- [快速开始](#快速开始)
+- [快速上手](#快速上手)
 - [演示](#演示)
 - [文档](#文档)
 - [贡献](#贡献)
 - [社区](#社区)
@@ -55,45 +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. 预处理数据**
 创建 `pretrain.json`（`seq` 策略的预处理配置）：
 ```json
 {
    "version": 1,
    "input": {"sections": [{"field": "text", "action": "train"}]},
    "preprocessing": {"max_seq_len": 2048},
    "output": {"storage_format": "bin"}
 }
 ```
 ```bash
-CUDA_VISIBLE_DEVICES=0,1,2,3 python scripts/tools/train.py \
+python scripts/tools/preprocess.py data/*.jsonl -o output/ -c pretrain.json
    --train_type seq \
    --data_root_path /path/to/dataset \
    --param_path /path/to/model \
    --batch_size 4 \
    --accumulation_steps 8 \
    --max_lr 3e-4 \
    --warmup_steps 1000 \
    --n_epoch 1
 ```
-完整参数列表见[参数说明](./params.md)。
+**4. 训练**
 ```bash
 export CUDA_VISIBLE_DEVICES=0,1,2,3
 nohup python scripts/tools/train.py \
    --nprocs=4 \
    --parallel_mode=ddp \
    --train_type=seq \
    --data_root_path=/path/to/dataset \
    --param_path=/path/to/model \
    --batch_per_device=4 \
    --grad_accum_steps=8 \
    --warmup_ratio=0.05 \
    --max_lr=1e-4 \
    --max_grad_norm=1.0 \
    --weight_decay=0.1 \
    --window_size=2048 \
    --ckpt_interval=10000 \
    --ckpt_dir=./checkpoint \
    --random_seed=3407 \
    --label_smoothing=0.05 \
    > out.log 2> err.log &
 ```
 **5. 启动服务并调用**
 ```bash
 # 终端 1：启动服务
 python scripts/tools/server.py --param_path ./params --device cuda
 # 终端 2：发起请求
 curl http://localhost:8000/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{"messages":[{"role":"user","content":"你好"}],"max_tokens":512}'
 ```
 ### 演示
 查看 `scripts/demo/` 文件夹中的演示：
 ```bash
 # 下载模型权重（运行演示前必需）
 python scripts/demo/download.py                      # model → params/
 # 交互式流式聊天（多轮对话，保持历史记录）
 python scripts/demo/stream_chat.py
 # 在 >> 后输入消息，输入 !exit 退出
 # 批量生成（5 条硬编码提示词，非流式）
 python scripts/demo/generate_batch.py
 # 单条提示词自回归流式生成
 python scripts/demo/generate_ar.py
 ```
 所有生成演示默认使用 `temperature=0.8`、`top_p=0.95`、`top_k=50`、`max_tokens=2048`，需要 `params/` 目录包含模型权重（请先运行 `download.py`）。
 观看 [bilibili](https://www.bilibili.com/video/BV1fuLB6yEj6) 上的视频演示。
 ---
 更多选项请参考[文档](#文档)。
 #### 文本生成
 从 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
@@ -107,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
@@ -126,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 驱动数据预处理 |
 ### 贡献
@@ -223,7 +247,7 @@ python scripts/demo/generate_ar.py
 - **GitHub Issues**: [问题追踪](https://github.com/ViperEkura/AstrAI/issues)
 - **Discussions**: [GitHub 讨论区](https://github.com/ViperEkura/AstrAI/discussions)
- **HuggingFace**: [模型中心](https://huggingface.co/ViperEk)
+- **HuggingFace**: [模型中心](https://huggingface.co/ViperEkura)
 ### 许可证
@@ -1,237 +1,130 @@
-# 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`): HDF5 data loading, checkpoint management
+                                              ↓
-
+                                      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. Serialization (`astrai/serialization.py`)
+Raw text is tokenized via `AutoTokenizer.encode()` and saved as HDF5 (`.h5`) or binary (`.bin` + `meta.json`) files with keyed tensor groups.
- **`save_h5`**: Saves tensors by groups as HDF5 files (`.h5`), each key maps to a list of tensors
+### 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 cos/sin cache
 - **`RMSNorm`**: Layer normalization
 ### 4. Training Module
 #### 4.1 Training Context (`train_context.py`)
 - **`TrainContext`**: Dataclass holding model, optimizer, dataloader, strategy, scheduler, checkpoint state
 - **`TrainContextBuilder`**: Builder pattern — takes checkpoint for resume, builds all components
 #### 4.2 Trainer (`trainer.py`)
 The training loop is nested: **epoch** → **batch** (with step phase interspersed):
 ```
 on_train_begin
  on_epoch_begin
    for each batch:
      if iteration % accumulation_steps == 0:        ← step phase
        on_step_begin → optimizer.step() → zero_grad → on_step_end
                                                      ← batch phase
      on_batch_begin → strategy(batch) → loss → backward → on_batch_end
      iteration += 1
    on_epoch_end
 on_train_end
 ```
-Key points:
+The output `meta.json` records the storage format, key names, dtype, total token count, and tensor shapes for each shard.
 - `on_step_*` wraps optimizer step (fires every `accumulation_steps` batches)
 - `on_batch_*` wraps loss computation (fires every batch)
 - `SchedulerCallback` fires on `on_batch_end` — LR scheduler steps every batch
 - `GradientClippingCallback` fires on `on_step_begin`
-#### 4.3 Strategy (`strategy.py`)
+### 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_begin`
 - **`SchedulerCallback`**: `scheduler.step()` on `on_batch_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, page_table, status (PENDING/RUNNING/FINISHED/ABORTED)
+All three inherit `Store` (base, owns `_data`/`_cum`/`_offsets`/`_normalize`) plus the `Streamable` and `Recordable` mixins, so every backend supports both `fetch(begin, end, keys)` (stream) and `fetch_record(index, keys)` (record) APIs.
 - **`PagedCache`**: Bitmask-based page allocator with page-table-indirected read/write
 - **`CacheView`**: Batch view bundling cache + page table for attention layers
 - **`sample()`**: Temperature → top-k → top-p → multinomial
-#### 5.3 Server (`server.py`)
+**H5Store**: Reads HDF5 files. Tensors are loaded into host memory and normalized into segmented storage. `segments_are_records=True` — each `data_i` dataset is one record.
 - FastAPI with OpenAI `/v1/chat/completions` and Anthropic `/v1/messages` endpoints
 - Streaming via SSE, health check at `/health`, stats at `/stats`
-### 6. Tokenizer Module
+**MmapStore**: Memory-maps `.bin` files. OS page cache sharing is native — no explicit `share_memory_()` needed. Uses `torch.from_numpy(np.memmap(...))`. `segments_are_records=False` — bin segments are contiguous streams; record access is driven by `_offsets` (written when `save_bin(..., record_keys=...)` was used at preprocessing time).
- **`AutoTokenizer`**: Wraps HuggingFace tokenizers (BBPE); `encode`/`decode`/`apply_chat_template`
+**JsonlStore**: On-the-fly tokenization of raw JSONL files at load time. Requires a `dataset_config.json` alongside the `.jsonl` files following the same `PipelineConfig` schema with an additional `tokenizer_path` field. Two modes: eager (default, applies `TokenizeTransform` to all records at load) and lazy (`processor=fn` given, defers tokenisation to `fetch_record` — used by DPO/GRPO).
 - **`ChatTemplate`**: Jinja2-based template rendering for multi-turn chat
-### 7. Factory & Parallel
+All backends normalise tensors into `Store._data[Dict[str, List[Tensor]]]` + `Store._cum[Dict[str, List[int]]]` (cumulative lengths for bisect-based stream indexing) + `Store._offsets[Dict[str, List[int]]]` (per-record offsets for record-mode indexing). Nested keys (GRPO `responses`/`masks` as `List[List[Tensor]]`) are stored as-is and excluded from both bookkeepings — they are only accessed record-by-record.
- **`Registry` / `BaseFactory`**: Decorator-based component registration
+## Data Keys by Training Type
 - **`spawn_parallel_fn`**: Multi-process DDP launcher with NCCL backend
 - **`ParallelModel` / `ColumnParallelLinear` / `RowParallelLinear`**: Tensor model parallelism
-## Training Data Flow — Detailed Steps
+| Type | Storage Keys | Access Mode |
 |------|-------------|-------------|
 | `seq` | `sequence` (→ input_ids, target_ids via offset-by-1) | stream (`fetch`) |
 | `sft` | `sequence`, `loss_mask`, `position_ids` | stream (`fetch`) |
 | `dpo` | `chosen`, `rejected`, `chosen_mask`, `rejected_mask` | record (`fetch_record`) |
 | `grpo` | `prompts`, `responses`, `masks`, `rewards` | record (`fetch_record`) |
-1. **Data Preparation**
+## Dataset Architecture
   - Raw text → token IDs via `AutoTokenizer.encode()`
   - Save as `.h5` files (groups of tensor lists per data key)
-2. **Dataset Loading**
+```
-   - `BaseDataset.load()` calls `load_h5()`, builds `MultiSegmentFetcher`
+DatasetFactory.load(train_type, load_path, window_size, stride=None,
-   - Sliding window of `window_size` with `stride` determines sample boundaries
+                    storage_type=None, tokenizer_path=None,
                    max_len=2048, store=None)
  → BaseDataset.load(load_path, storage_type=None)
    → detect_format(load_path)
    → StoreFactory.create(storage_type)
    → Store.load(load_path)
      → _normalize(raw)  # base Store, shared by both backends
        → Store._data[Dict[str, List[Tensor]]]
          + _cum[Dict[str, List[int]]]   (stream mode)
          + _offsets[Dict[str, List[int]]]  (record mode)
-3. **Sampling & Batching**
+Stream datasets (SEQ/SFT):
-   - `ResumableDistributedSampler` produces shuffled index sequences
+  BaseDataset.__getitem__(idx)
-   - `DataLoader` fetches `[batch_size, window_size]` tensors via `__getitem__`
+    → get_index(idx) → [begin, end)
    → Store.fetch(begin, end, keys) → Tensor / Dict[str, Tensor]
-4. **Strategy Forward**
+Record datasets (DPO/GRPO via RecordDataset):
-   - Strategy receives batch, calls `Transformer.forward()` for logits
+  RecordDataset.__getitem__(idx)
-   - Computes task-specific loss (cross-entropy, DPO, GRPO)
+    → Store.fetch_record(idx, keys) → Tensor / Dict[str, Tensor]
 ```
-5. **Backward & Accumulation**
+Class hierarchy: `BaseDataset` ← `SEQDataset` / `SFTDataset` (stream); `BaseDataset` ← `RecordDataset` ← `DPODataset` / `GRPODataset` (record).
   - `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**
+`window_size` = max input length, `stride` = step between consecutive samples (defaults to `window_size`, optional). Only meaningful for stream datasets — record datasets ignore both. `storage_type` defaults to `None` (auto-detect via `detect_format`).
   - `CheckpointCallback` saves `model.state_dict()` + metadata to safetensors at `ckpt_interval` iterations
   - Does NOT save optimizer/scheduler state (resume resets those)
-## Inference Data Flow — Detailed Steps
+`tokenizer_path` triggers lazy on-the-fly tokenisation for record datasets on raw JSONL (DPO builds a `dpo_processor`; SEQ/SFT/pre-tokenised backends ignore it). `store` (pre-built `Store`) bypasses `load_path`/`storage_type`/`tokenizer_path` entirely — the caller controls Store construction.
-1. **Model Loading**
+`Store.fetch(begin, end, keys)` (stream mode, on `Streamable`): accepts a single key (`str`) returning a `Tensor`, or a list of keys returning `Dict[str, Tensor]`. Internally uses `bisect` across multi-segment tensors. Raises `RuntimeError("Store not loaded")` if called before `load()`.
   - `AutoModel.from_pretrained(path)` loads weights from safetensors
   - `torch.inference_mode()` wraps generation
-2. **Prompt Construction**
+`Store.fetch_record(index, keys)` (record mode, on `Recordable`): same key API. Uses `_offsets[key]` when present (bin layout with per-record offsets), otherwise indexes `_data[key]` directly (H5/JSONL where each segment is one record).
   - Messages → `apply_chat_template(messages, tokenize=False)` → prompt string
   - `tokenizer.encode(prompt)` → token IDs (truncated to `max_prompt_len`)
-3. **Continuous Batching Loop**
+## Sampler
   - **Cleanup**: Finished tasks → `stream_callback(STOP)`, free KV pages
   - **Refill**: Pop from waiting queue, `PagedCache.alloc_n()` for prompt pages
   - **Prefill**: Group by prompt length, run full forward with `start_pos=0`
   - **Decode**: Pick position group with most tasks, single-token forward:
     - Model forward → `logits` → `sample()` → next token ID
     - Append to `output_ids`, update `output_tokens`
     - `_maybe_alloc_page()` grows page table as needed
     - `stream_callback(token)` for streaming clients
-4. **Output**
+`ResumableDistributedSampler` supports checkpoint-aware distributed sampling:
   - `tokenizer.decode(output_ids)` → text
   - Return to caller (streaming: token-by-token; non-streaming: complete string)
-## Checkpoint & Serialization
+- Tracks `start_epoch` / `start_iter` for resume
 - Shuffle via `torch.Generator(seed + epoch)`
 - Per-replica index slicing for DDP
- **Training Checkpoint**: safetensors weights + epoch/iteration metadata. Optimizer/scheduler state is NOT persisted.
+## DataLoader
 - **Inference Loading**: `AutoModel.from_pretrained()` loads from the same safetensors format.
 - **Dataset Serialization**: HDF5 with shared memory support for large-scale pre-training data.
-> Document Update Time: 2026-05-09
+Standard PyTorch `DataLoader` with configurable `batch_size`, `num_workers`, `pin_memory`, `prefetch_factor`. Sampler produces indices; dataloader fetches tensor batches via `__getitem__`.
 > Document Update Time: 2026-07-19
@@ -1,719 +0,0 @@
 ## 1. Why I Created This Project
 There are many large language models on the market today, such as GPT, LLaMA, and others, with tens of billions or even hundreds of billions of parameters. But honestly, these models have extremely high hardware requirements, making them inaccessible for ordinary developers. I thought: **Can we create a model that is both useful and can run on ordinary computers?** This is also what most people currently hope for - a locally deployable AI project that achieves complete privatization while maintaining some level of intelligence.
 Thus, the AstrAI project was born - 1B parameters, Chinese-English bilingual, supporting dialogue, text generation, and the training code is open source!
 ## 2. System Architecture
 ```mermaid
 classDiagram
    namespace config {
        class ModelConfig {
            +int vocab_size
            +int dim
            +int n_layers
            +float norm_eps
            +int dim_ffn
            +bool tie_weight
            +int max_len
            +float rope_theta
            +int n_heads
            +int n_kv_heads
            +bool use_qk_norm
            +bool use_gated_attention
            +load(config_path) ModelConfig
            +save(config_path)
        }
        class TrainConfig {
            +nn.Module model
            +str strategy
            +Dataset dataset
            +Callable optimizer_fn
            +Callable scheduler_fn
            +int n_epoch
            +int batch_size
            +int accumulation_steps
            +float max_grad_norm
            +int start_epoch
            +int start_batch
            +str ckpt_dir
            +int ckpt_interval
            +int random_seed
            +int num_workers
            +int prefetch_factor
            +bool pin_memory
            +int nprocs
            +str backend
            +str master_addr
            +str master_port
            +Callable parallel_wrapper
            +Callable state_dict_fn
            +str device_type
            +dict extra_kwargs
            +validate()
        }
    }
    namespace dataset {
        class BaseDataset {
            +int window_size
            +int stride
            +MultiSegmentFetcher fetcher
            +load(load_path)
            +__getitem__(index)
            +__len__()
        }
        class SEQDataset {
            +__getitem__(index) Dict
        }
        class SFTDataset {
            +__getitem__(index) Dict
        }
        class DPODataset {
            +__getitem__(index) Dict
        }
        class GRPODataset {
            +__getitem__(index) Dict
        }
        class BaseSegmentFetcher {
            +List[Tensor] segments
            +List[int] cum_lengths
            +int total_length
            +fetch_data(begin_idx, end_idx) Tensor
        }
        class MultiSegmentFetcher {
            +Dict multi_fetchers
            +List multi_keys
            +key_fetch(begin_idx, end_idx, keys) Dict
            +fetch_data(begin_idx, end_idx) Dict
        }
        class ResumableDistributedSampler {
            +int epoch
            +int iter
        }
        class DatasetFactory {
            +Registry _registry
            +register(name) decorator
            +create(train_type, window_size, stride) BaseDataset
            +load(train_type, load_path, window_size, stride) BaseDataset
        }
    }
    namespace serialization {
        class Checkpoint {
            +dict state_dict
            +int epoch
            +int iteration
            +save(save_dir)
            +load(save_dir) Checkpoint
        }
    }
    namespace model {
        class AutoModel {
            +ModelConfig config
            +Registry _registry
            +register(model_type) decorator
            +get_model_class(model_type) Type
            +from_pretrained(path, disable_random_init) nn.Module
            +save_pretrained(save_directory)
            +to(*args, **kwargs) Self
        }
        class Transformer {
            +ModelConfig config
            +RotaryEmbedding rotary_embedding
            +Embedding embed_tokens
            +ModuleList layers
            +RMSNorm norm
            +Linear lm_head
            +forward(input_ids, input_mask, paged_cache, start_pos) Dict
            +load_state_dict(state_dict)
            +state_dict()
        }
        class DecoderBlock {
            +GQA attention
            +RMSNorm input_norm
            +MLP mlp
            +RMSNorm post_attention_norm
            +forward(x, rotary_emb, attention_mask, paged_cache, start_pos) Tensor
        }
        class GQA {
            +int n_heads
            +int n_kv_heads
            +int head_dim
            +Linear q_proj, k_proj, v_proj, o_proj
            +RMSNorm q_norm, k_norm
            +forward(x, rotary_emb, mask, paged_cache, start_pos) Tensor
        }
        class MLA {
            +int n_heads
            +int n_kv_heads
            +int head_dim
            +int kv_lora_rank
            +int qk_nope_head_dim
            +int qk_rope_head_dim
            +Linear q_proj, kv_a_proj, kv_b_proj
            +Linear o_proj
            +RMSNorm kv_norm
            +forward(x, rotary_emb, mask, paged_cache, start_pos) Tensor
        }
        class MLP {
            +Linear up, gate, down
            +forward(x) Tensor
        }
        class RMSNorm {
            +Parameter weight
            +float norm_eps
            +forward(x) Tensor
        }
        class Linear {
            +Parameter weight
            +Parameter bias
            +forward(x) Tensor
        }
        class RotaryEmbedding {
            +int dim
            +int max_len
            +float base
            +forward(x, start_pos) Tuple[Tensor, Tensor]
        }
        class Embedding {
            +Parameter weight
            +forward(x) Tensor
        }
    }
    namespace tokenize {
        class AutoTokenizer {
            +List[int] stop_ids
            +int bos_id
            +int eos_id
            +int pad_id
            +vocab_size int
            +encode(tokens, out_ids, add_special_tokens) List[int]
            +decode(tokens, skip_special_tokens) str
            +apply_chat_template(messages, tokenize) Union[str, List[int]]
            +set_chat_template(template)
            +load(path)
            +from_pretrained(path) AutoTokenizer
            +save_pretrained(save_path)
        }
        class ChatTemplate {
            +String template_str
            +render(messages, system_prompt, **extra_variables) str
            +from_string(template) ChatTemplate
        }
    }
    namespace factory {
        class Registry {
            +Dict _entries
            +register(name, component_cls, category, priority)
            +get(name) Type
            +list_names() List[str]
        }
        class BaseFactory {
            +Registry _registry
            +register(name, category, priority) decorator
            +create(name, *args, **kwargs) T
            +list_registered() list
        }
    }
    namespace trainer {
        class Trainer {
            +TrainConfig train_config
            +List[TrainCallback] callbacks
            +train(checkpoint)
            +_build_context(checkpoint) TrainContext
            +_get_default_callbacks() List[TrainCallback]
        }
        class TrainContext {
            +nn.Module model
            +BaseStrategy strategy
            +DataLoader dataloader
            +Optimizer optimizer
            +LRScheduler scheduler
            +Checkpoint checkpoint
            +int epoch
            +int iteration
            +float loss
            +int world_size
            +int rank
        }
        class TrainContextBuilder {
            +TrainConfig config
            +with_checkpoint(checkpoint) TrainContextBuilder
            +build() TrainContext
        }
        class BaseStrategy {
            +nn.Module model
            +str device
            +compute_loss(batch) Tensor
        }
        class StrategyFactory {
            +Registry _registry
            +register(name) decorator
            +create(model, train_type, device, **kwargs) BaseStrategy
        }
        class SEQStrategy {
            +float label_smoothing
            +compute_loss(batch) Tensor
        }
        class SFTStrategy {
            +float label_smoothing
            +compute_loss(batch) Tensor
        }
        class DPOStrategy {
            +nn.Module ref_model
            +float beta
            +str reduction
            +compute_loss(batch) Tensor
        }
        class GRPOStrategy {
            +nn.Module ref_model
            +float clip_eps
            +float kl_coef
            +int group_size
            +compute_loss(batch) Tensor
        }
        class BaseScheduler {
            +get_lr() List[float]
            +step()
        }
        class SchedulerFactory {
            +Registry _registry
            +register(name) decorator
            +create(optimizer, schedule_type, **kwargs) BaseScheduler
        }
        class CosineScheduler {
            +int warmup_steps
            +int lr_decay_steps
            +float min_rate
        }
        class SGDRScheduler {
            +int warmup_steps
            +int cycle_length
            +float min_rate
            +int t_mult
        }
        class TrainCallback {
            +on_train_begin(context)
            +on_train_end(context)
            +on_epoch_begin(context)
            +on_epoch_end(context)
            +on_step_begin(context)
            +on_step_end(context)
            +on_batch_begin(context)
            +on_batch_end(context)
            +on_error(context)
        }
        class GradientClippingCallback {
            +float max_grad_norm
            +on_step_begin(context)
        }
        class SchedulerCallback {
            +on_train_begin(context)
            +on_batch_end(context)
        }
        class CheckpointCallback {
            +str save_dir
            +int interval
            +_save_checkpoint(context)
            +on_batch_end(context)
            +on_train_end(context)
            +on_error(context)
        }
        class ProgressBarCallback {
            +int num_epoch
            +on_epoch_begin(context)
            +on_batch_end(context)
            +on_epoch_end(context)
        }
        class MetricLoggerCallback {
            +str log_dir
            +int save_interval
            +on_batch_end(context)
            +on_train_end(context)
        }
        class CallbackFactory {
            +Registry _registry
            +register(name) decorator
            +create(name, **kwargs) TrainCallback
        }
    }
    namespace inference {
        class InferenceEngine {
            +nn.Module model
            +AutoTokenizer tokenizer
            +InferenceScheduler scheduler
            +int max_batch_size
            +Optional int max_seq_len
            +generate(prompt, stream, max_tokens, temperature, top_p, top_k) Union[Generator, str, List[str]]
            +generate_with_request(request) Union[Generator, str, List[str]]
            +generate_async(prompt, max_tokens, temperature, top_p, top_k) AsyncGenerator
            +get_stats() Dict
            +shutdown()
        }
        class InferenceScheduler {
            +nn.Module model
            +AutoTokenizer tokenizer
            +PagedCache page_cache
            +int max_batch_size
            +int max_seq_len
            +int max_prompt_len
            +int page_size
            +List waiting_queue
            +List active_tasks
            +add_task(prompt, max_tokens, temperature, top_p, top_k, stream_callback) str
            +remove_task(task_id)
            +start()
            +stop()
            +get_stats() Dict
        }
        class PagedCache {
            +int page_size
            +int _free_mask
            +List[int] _refs
            +Tensor k_cache
            +Tensor v_cache
            +alloc() int
            +alloc_n(n) List[int]
            +free(idx)
            +bind(page_table, total_len) CacheView
            +write(layer_id, page_table, start_pos, k, v)
            +gather(layer_id, page_table) Tuple[Tensor, Tensor]
        }
        class CacheView {
            +PagedCache _cache
            +Tensor _page_table
            +int _total_len
            +write(layer_id, start_pos, k, v)
            +gather(layer_id) Tuple[Tensor, Tensor]
        }
        class Task {
            +str task_id
            +List prompt_ids
            +int max_tokens
            +float temperature
            +float top_p
            +int top_k
            +TaskStatus status
            +List output_ids
            +int input_tokens
            +int output_tokens
            +List[int] page_table
            +int n_pages
            +float arrival_time
            +float finish_time
            +Callable stream_callback
            +int next_pos
            +is_finished(stop_ids) bool
        }
        class TaskStatus {
            <<enumeration>>
            PENDING
            RUNNING
            FINISHED
            ABORTED
        }
        class GenerationRequest {
            +List[Dict] messages
            +GenerationParams params
            +bool stream
        }
        class GenerationParams {
            <<value object>>
            +int top_k
            +float top_p
            +float temperature
            +int max_tokens
        }
        class BaseSamplingStrategy {
            <<abstract>>
            +apply(logits, filter_value) Tensor
        }
        class TemperatureStrategy {
            +float temperature
            +apply(logits, filter_value) Tensor
        }
        class TopKStrategy {
            +int top_k
            +apply(logits, filter_value) Tensor
        }
        class TopPStrategy {
            +float top_p
            +apply(logits, filter_value) Tensor
        }
        class SamplingPipeline {
            +List strategies
            +apply(logits, filter_value) Tensor
            +sample(logits, filter_value) Tensor
        }
        class _Result {
            +List[str] tokens
            +List[str] results
            +List[bool] _done
            +append(token, idx)
            +get_results() List[str]
            +pop_all() List[str]
            +wait(timeout) bool
        }
        class ChatMessage {
            +str role
            +str content
        }
        class ChatCompletionRequest {
            +List[ChatMessage] messages
            +float temperature
            +float top_p
            +int top_k
            +int max_tokens
            +bool stream
            +Optional[str] stop
            +Optional[int] n
        }
    }
    namespace parallel {
        class ParallelFunctions {
            +spawn_parallel_fn(fn, nprocs)
            +setup_parallel(rank, world_size, backend, master_addr, master_port, device_type)
        }
        class ParallelModel {
            +dist.ProcessGroup process_group
            +int rank
            +int world_size
        }
        class ColumnParallelLinear {
            +forward(x) Tensor
        }
        class RowParallelLinear {
            +forward(x) Tensor
        }
    }
    %% Relationships
    TrainConfig --> ModelConfig : uses
    TrainConfig --> BaseDataset : uses
    TrainConfig --> StrategyFactory : selects
    StrategyFactory ..> BaseStrategy : creates
    BaseStrategy <|-- SEQStrategy
    BaseStrategy <|-- SFTStrategy
    BaseStrategy <|-- DPOStrategy
    BaseStrategy <|-- GRPOStrategy
    DPOStrategy --> Transformer : uses
    GRPOStrategy --> Transformer : uses
    Trainer --> TrainConfig : configures
    Trainer --> TrainContextBuilder : builds
    Trainer --> TrainCallback : manages
    TrainContextBuilder --> TrainContext : creates
    Checkpoint ..> Checkpoint : saves/loads
    TrainContext --> Checkpoint : manages
    TrainContext --> BaseStrategy : uses
    TrainContext --> BaseScheduler : uses
    SchedulerFactory ..> BaseScheduler : creates
    BaseScheduler <|-- CosineScheduler
    BaseScheduler <|-- SGDRScheduler
    CallbackFactory ..> TrainCallback : creates
    TrainCallback <|-- GradientClippingCallback
    TrainCallback <|-- SchedulerCallback
    TrainCallback <|-- CheckpointCallback
    TrainCallback <|-- ProgressBarCallback
    TrainCallback <|-- MetricLoggerCallback
    InferenceEngine --> InferenceScheduler : uses
    InferenceEngine --> GenerationRequest : uses
    GenerationRequest --> GenerationParams : contains
    InferenceScheduler --> Task : manages
    Task --> TaskStatus : uses
    InferenceScheduler --> TaskStatus : uses
    InferenceScheduler --> PagedCache : uses
    InferenceScheduler --> Transformer : uses
    InferenceEngine --> Transformer : uses
    InferenceEngine --> _Result : uses
    BaseSamplingStrategy <|-- TemperatureStrategy
    BaseSamplingStrategy <|-- TopKStrategy
    BaseSamplingStrategy <|-- TopPStrategy
    SamplingPipeline --> BaseSamplingStrategy : composes
    BaseDataset <|-- SEQDataset
    BaseDataset <|-- SFTDataset
    BaseDataset <|-- DPODataset
    BaseDataset <|-- GRPODataset
    DatasetFactory ..> BaseDataset : creates
    MultiSegmentFetcher --> BaseSegmentFetcher : uses
    BaseDataset --> MultiSegmentFetcher : uses
    AutoModel <|-- Transformer
    AutoModel --> ModelConfig : contains
    Transformer --> DecoderBlock : uses
    Transformer --> RotaryEmbedding : uses
    Transformer --> Embedding : uses
    DecoderBlock --> GQA : uses
    DecoderBlock --> MLP : uses
    DecoderBlock --> RMSNorm : uses
    TrainContextBuilder --> ResumableDistributedSampler : creates
    ResumableDistributedSampler --> BaseDataset : samples
    ParallelModel <|-- RowParallelLinear
    ParallelModel <|-- ColumnParallelLinear
    AutoTokenizer --> ChatTemplate : uses
    TrainConfig --> DatasetFactory : selects
    TrainConfig --> SchedulerFactory : selects
    TrainConfig --> CallbackFactory : selects
    AutoModel ..> AutoTokenizer : loads with
    BaseFactory <|-- DatasetFactory
    BaseFactory <|-- StrategyFactory
    BaseFactory <|-- SchedulerFactory
    BaseFactory <|-- CallbackFactory
 ```
 ### Module Overview
 | Module | Components | Description |
 |--------|------------|-------------|
 | **astrai.config** | ModelConfig, TrainConfig | Configuration management |
 | **astrai.dataset** | BaseDataset, SEQDataset, SFTDataset, DPODataset, GRPODataset, BaseSegmentFetcher, MultiSegmentFetcher, ResumableDistributedSampler, DatasetFactory | Dataset loading and management |
 | **astrai.serialization** | Checkpoint, save_h5, load_h5 | Model serialization and checkpoint management |
 | **astrai.model** | AutoModel, Transformer, DecoderBlock, GQA, MLA, MLP, RMSNorm, Linear, RotaryEmbedding, Embedding | Neural network model |
 | **astrai.tokenize** | AutoTokenizer, ChatTemplate | Tokenizer and chat template |
 | **astrai.trainer** | Trainer, TrainContext, TrainContextBuilder, BaseStrategy, StrategyFactory, BaseScheduler, SchedulerFactory, TrainCallback, CallbackFactory | Training workflow management |
 | **astrai.inference** | InferenceEngine, InferenceScheduler, PagedCache, CacheView, Task, TaskStatus, GenerationParams, GenerationRequest, BaseSamplingStrategy, TemperatureStrategy, TopKStrategy, TopPStrategy, SamplingPipeline, ChatMessage, ChatCompletionRequest | Inference service with continuous batching and paged KV cache |
 | **astrai.parallel** | ParallelFunctions, ParallelModel, ColumnParallelLinear, RowParallelLinear | Distributed parallel |
 | **astrai.factory** | Registry, BaseFactory | Generic component registration |
 ### Design Patterns
 | Pattern | Classes | Purpose |
 |---------|---------|---------|
 | **Strategy** | `BaseStrategy`, `SEQStrategy`, `SFTStrategy`, `DPOStrategy`, `GRPOStrategy`, `StrategyFactory` | Flexible training strategy switching, supports SEQ/SFT/DPO/GRPO |
 | **Builder** | `TrainContextBuilder` | Chain-building training context, step-by-step initialization of components |
 | **Factory** | `StrategyFactory`, `SchedulerFactory`, `DatasetFactory`, `CallbackFactory`, `BaseFactory` | Decorator registration mechanism, dynamically create training strategies, schedulers, datasets, and callbacks |
 | **Observer** | `TrainCallback`, `CallbackFactory` | Callback mechanism for training process monitoring (checkpoint, early stopping, metrics) |
 | **Context** | `TrainContext` | Training process state container with model, optimizer, scheduler and checkpoint |
 | **Registry** | `BaseFactory`, `Registry` | Generic component registration with category and priority support |
 | **Object Pool** | `PagedCache` | Page-based KV cache with O(1) alloc/free via bitmask |
 | **Strategy (Sampling)** | `BaseSamplingStrategy`, `TemperatureStrategy`, `TopKStrategy`, `TopPStrategy`, `SamplingPipeline` | Composable logit transformations with temperature, top-k, top-p |
 | **Producer-Consumer** | `InferenceScheduler`, `Task`, `waiting_queue`, `active_tasks` | Continuous batching with dynamic task queue management |
 | **Event-Driven** | `threading.Event`, `_task_event` | Non-blocking wait mechanism for task scheduling using Python's `threading` module |
 | **AutoModel Registry** | `AutoModel`, `Transformer` | Model type registration and dynamic loading via decorator pattern |
 | **Generator Pattern** | `_Result`, `GenerationRequest` | Event-based result notification for streaming/non-streaming generation |
 ### Core Relationships
 1. **Configuration → Training**: `TrainConfig` contains `ModelConfig`, holds model, dataset, optimizer and other references
 2. **Training Flow**: `Trainer` → `TrainContextBuilder` → `TrainContext`, uses `BaseStrategy` to compute loss
 3. **Strategy Selection**: `StrategyFactory` creates corresponding strategy instance based on `train_type`
 4. **Inference Flow**: `InferenceEngine` → `InferenceScheduler` → `Transformer`, uses `PagedCache` for paged KV cache management and `SamplingPipeline` for efficient continuous batching with streaming/non-streaming
 5. **Distributed Support**: `spawn_parallel_fn` and `setup_parallel` provide multi-process training capability for `Trainer`
 6. **Dataset Loading**: `DatasetFactory` creates datasets (SEQDataset, SFTDataset, DPODataset, GRPODataset), supports HDF5 loading via `BaseSegmentFetcher` and `MultiSegmentFetcher`
 7. **Checkpoint Management**: `Checkpoint` handles model state serialization/deserialization with safetensors
 8. **Scheduler Support**: `SchedulerFactory` creates learning rate schedulers (CosineScheduler, SGDRScheduler)
 9. **AutoModel Loading**: `AutoModel.from_pretrained()` dynamically loads model based on `config.json` model_type, uses `Registry` pattern for model type registration
 ## 3. Training Process
 The common training process for large language models (LLM) typically includes three stages: **Pre-training (SEQ)**, **Supervised Fine-Tuning (SFT)**, and **Reinforcement Learning from Human Feedback (DPO/GRPO)**. This system is designed to support seamless end-to-end flow, achieving efficient switching and state management of different training stages through modular strategies.
 ### Core Formulas
 **Pre-training (SEQ):**
 $$
 L_{\text{PT}} = - \sum_{t=1}^{T} \log P(x_t \mid x_{\lt t}; \theta)
 $$
 **SFT:**
 $$
 L_{\text{SFT}} = - \sum_{t=P+1}^{P+L} \log P(s_t \mid s_{\lt t}; \theta)
 $$
 **DPO:**
 $$
 L_{\text{DPO}} = -\mathbb{E}_{(x, y_w, y_l) \sim D} \left[ \log \sigma\left( \beta \log \frac{\pi_\theta(y_w \mid x)}{\pi_{\text{ref}}(y_w \mid x)} - \beta \log \frac{\pi_\theta(y_l \mid x)}{\pi_{\text{ref}}(y_l \mid x)} \right) \right]
 $$
 **GRPO:**
 GRPO (Group Relative Policy Optimization) computes advantages from multiple responses to the same prompt, then optimizes using a PPO-style clipped objective:
 $$
 \text{Advantage}_i = \frac{r_i - \mu}{\sigma + \epsilon}
 $$
 Where $r_i$ is the reward for the $i$-th response, $\mu$ and $\sigma$ are the mean and standard deviation of group rewards.
 $$
 L_{\text{GRPO}} = -\mathbb{E} \left[ \min\left( \frac{\pi_\theta(a|s)}{\pi_{\text{ref}}(a|s)} \cdot A, \text{clip}\left(\frac{\pi_\theta(a|s)}{\pi_{\text{ref}}(a|s)}, 1-\epsilon, 1+\epsilon\right) \cdot A \right) \right] + \lambda \cdot D_{KL}
 $$
 The KL divergence term uses mean squared error approximation:
 $$
 L_{KL} = \lambda \cdot \mathbb{E} \left[ (\log \pi_\theta - \log \pi_{\text{ref}})^2 \right]
 $$
 The final loss is the sum of both: $L = L_{\text{policy}} + L_{KL}$
 Through the above three-stage progressive training, the model completes its evolution from a general language foundation to a specialized, highly-aligned dialogue intelligence.
 > Document Update Time: 2026-04-09
@@ -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 24 layers of Transformer blocks, with 1.0 billion parameters. Transformer is an autoregressive model that calculates the relationship between all previous tokens to obtain the probability distribution of the next token.
 The model now uses the **AutoModel** base class for flexible loading and saving:
 ```python
 from astrai.model import AutoModel
 # Load model from checkpoint
 model = AutoModel.from_pretrained("path/to/model")
 # Save model to new directory
 model.save_pretrained("path/to/save")
 ```
 The Transformer model is registered via `@AutoModel.register('transformer')` decorator, allowing easy extension for new model types.
 ```mermaid
 flowchart TB
    subgraph Layers["Transformer Layers"]
        direction TB
        A[Input Embedding] --> B[Transformer Block\nLayer 1]
        B --> C[Transformer Block\nLayer ...]
        C --> D[Transformer Block\nLayer 32]
        D --> E[RMSNorm]
        E --> F[Linear]
        F --> G[SoftMax]
    end
    subgraph TransformerBlock["Transformer Block"]
        direction TB
        H[x] --> I[RMSNorm]
        I --> J[Linear → Q/K/V]
        J --> K[Q]
        J --> L[K]
        J --> M[V]
        K --> N[RoPE]
        L --> O[RoPE]
        N --> P["Q @ K^T / sqrt(d)"]
        O --> P
        P --> Q[Masked SoftMax]
        Q --> R[S @ V]
        M --> R
        R --> S[Linear]
        S --> T[+]
        H --> T
        T --> U[RMSNorm]
        U --> V["Linear (gate)"]
        U --> W["Linear (up)"]
        V --> X[SiLU]
        X --> Y[×]
        W --> Y
        Y --> Z["Linear (down)"]
        Z --> AA[+]
        T --> AA
        AA --> BB[x']
    end
    classDef main fill:#e6f3ff,stroke:#0066cc;
    classDef block fill:#fff2e6,stroke:#cc6600;
    class Layers main;
    class TransformerBlock block;
 ```
 What is an autoregressive model? After splitting a sentence into tokens, the model predicts the probability distribution of the next token. This means the model calculates the probability of the next possible token and its corresponding probability based on the given context (the sequence of tokens that have already appeared).
 #### 1. Autoregression
 In autoregressive modeling, when a sentence is tokenized into a sequence of tokens, the model learns to predict what comes next. Given a sequence of tokens as input, the model calculates a probability distribution over all possible next tokens. This distribution tells us how likely each potential next token is, given the current context.
 For instance, if the input sequence contains tokens representing a question, the model might predict that certain response tokens have higher probabilities than others. The sampling process then selects one token from this distribution—controlled by parameters like top_k, top_p, and temperature—to serve as the next token in the sequence.
 Once a token is selected, it is appended to the input sequence, and the model repeats this process. The updated sequence is then fed back into the model to predict the next token. This iterative process continues until either a special end-of-sequence token is generated, or the maximum sequence length is reached. These control tokens are essential because without them, the model would continue generating tokens indefinitely, eventually exhausting available memory.
 #### 2. Causal Mask
 Transformers use attention mechanism. The input shape is generally [bsz, seq_len], and the output is [bsz, seq_len, n_dim]. To predict the next token, the model's input and output must be offset by one position. The target predicted by the model must be offset by one position, and during training we also use the offset-by-one method:
 ```
 sequence : [[1, 2, 3, 4, 5, 6]]
 input_ids: [[1, 2, 3, 4, 5]]
 target_ids: [[2, 3, 4, 5, 6]]
 ```
 The attention score calculation formula is:
 $$ s_{ij} = softmax(\frac{q_i^Tk_j}{\sqrt{d_k}}) $$
 $$ s_{ij} := s_{ij} + mask_{ij} $$
 Here, the attention score represents the degree to which the model attends to the similarity between two tokens.
 For decoder-only structure models, to prevent the model from "stealing" information from future positions, a mask needs to be added during attention calculation. We need to apply a mask before attention score calculation. This mask is typically a lower triangular matrix, and for a sequence of length n, its shape is [n, n]. Below is an example of how to create such a causal mask matrix for a sequence of length 5:
 ```
 [[0, -inf, -inf, -inf, -inf],
 [0,    0, -inf, -inf, -inf],
 [0,    0,    0, -inf, -inf],
 [0,    0,    0,    0, -inf],
 [0,    0,    0,    0,    0]]
 ```
 In this matrix, 0 represents positions that can be attended to, while -inf represents positions that should be masked (i.e., should not be attended to). Because this matrix ensures that after the softmax, the parts of the attention scores where $j > i$ change from `inf` to 0, meaning the model cannot see future information.
 #### 3. Rotary Position Embedding
 Rotary Position Embedding (RoPE) is a position encoding method designed to solve the problem of lacking direct modeling of sequence position information in Transformer models. Unlike traditional position encodings (such as sine and cosine function position encodings), RoPE embeds position information directly into the Query (Q) and Key (K) vectors, allowing the model to more naturally handle relative position relationships in sequences.
 $$ q_i = R_i W_q x_i $$
 $$ k_j = R_j W_k x_j $$
 $$ q_i^T k_j = (R_i W_q x_i)^T( R_j W_k x_j) = x_i^T W_q^T R_{i-j} W_k x_j $$
 The $R_{i-j}$ controls the attenuation of attention for different tokens at different relative distances. When the absolute value of $i - j$ is larger, the degree of attenuation is stronger. This approach allows the model to learn relative position relationships, enabling the model to scale and adapt to longer sequences.
 ## KV Cache Implementation
 According to the attention calculation formula:
 $$
 \begin{align*}
 o_i &= \sum_j s_{ij} v_{j} \newline
 s_{ij} &= \text{softmax}\left( \frac{q_{i} k_{j}}{\sqrt{d_k}} \right)
 \end{align*}
 $$
 Since the model is an autoregressive model, we only need to calculate for the last part of the sequence, meaning the index $i$ is fixed as the last element of the sequence, and we compute $o_{n}$:
 $$
 \begin{align*}
 o_n &= \sum_j s_{j}v_{j} \newline
 s_j &= \text{softmax}\left(\frac{q_n k_{j}}{\sqrt{d_k}} \right)
 \end{align*}
 $$
 If we expand the expression:
 $$
 o_n = \sum_j \text{softmax}\left(\frac{q_n k_{j}}{\sqrt{d_k}}\right)v_{j}
 $$
 In the above expression, only k and v have length indices, while $q$ does not. Therefore, during the calculation process, the input of $q$ is fixed as the last token from the previous input, while $k$ and $v$ need to be cached for parts of different lengths. Also, when caching, note that position encoding calculation should be performed before KV cache computation, otherwise there will be position encoding calculation errors.
 ### 4. AutoModel Loading
 The project now uses the **AutoModel** base class for flexible model loading and saving:
 ```python
 from astrai.model import AutoModel
 # Load model from checkpoint
 model = AutoModel.from_pretrained("path/to/model")
 # Save model to new directory
 model.save_pretrained("path/to/save")
 ```
 The Transformer model is registered via `@AutoModel.register('transformer')` decorator, allowing easy extension for new model types. The `from_pretrained` method automatically loads the `config.json` to determine the model type and uses safetensors format for weights.
 ### 5. Continuous Batching Inference
 The inference engine supports **continuous batching** for efficient batch processing:
 ```python
 from astrai.inference import InferenceEngine, GenerationRequest
 # Create inference engine with continuous batching
 engine = InferenceEngine(
    model=model,
    tokenizer=tokenizer,
 )
 # Use GenerationRequest with messages format
 request = GenerationRequest(
    messages=[
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": "Hello"},
    ],
    temperature=0.8,
    top_p=0.95,
    top_k=50,
    max_len=1024,
    stream=True,
 )
 # Generate with streaming
 for token in engine.generate_with_request(request):
    print(token, end="", flush=True)
 ```
 The continuous batching feature allows dynamic batch composition where new requests can join at any time and completed requests are released immediately.
 ## HTTP API Usage
 The inference server provides HTTP endpoints for remote inference. Start the server first:
 ```bash
 python -m scripts.tools.server --port 8000
 ```
 ### OpenAI-Compatible Endpoint
 The server provides an OpenAI-compatible chat completion endpoint at `/v1/chat/completions`:
 ```bash
 curl -X POST http://localhost:8000/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{
    "messages": [
      {"role": "system", "content": "You are a helpful assistant."},
      {"role": "user", "content": "Hello, how are you?"}
    ],
    "temperature": 0.8,
    "max_tokens": 2048,
    "stream": false
  }'
 ```
 **Request Parameters:**
 | Parameter | Type | Default | Description |
 |-----------|------|---------|-------------|
 | `messages` | List[dict] | Required | Chat messages with role and content |
 | `temperature` | float | 1.0 | Sampling temperature (0.0-2.0) |
 | `top_p` | float | 1.0 | Nucleus sampling threshold |
 | `top_k` | int | 50 | Top-k sampling parameter |
 | `max_tokens` | int | 1024 | Maximum tokens to generate |
 | `stream` | bool | false | Enable streaming response |
 **Response (non-streaming):**
 ```json
 {
  "id": "chatcmpl-1234567890",
  "object": "chat.completion",
  "created": 1234567890,
  "model": "astrai",
  "choices": [
    {
      "index": 0,
      "message": {"role": "assistant", "content": "Hello! I'm doing well..."},
      "finish_reason": "stop"
    }
  ],
  "usage": {
    "prompt_tokens": 20,
    "completion_tokens": 15,
    "total_tokens": 35
  }
 }
 ```
 ### Streaming Response
 Enable streaming for real-time token-by-token output:
 ```bash
 curl -X POST http://localhost:8000/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{
    "messages": [{"role": "user", "content": "Write a story"}],
    "stream": true,
    "max_tokens": 500
  }'
 ```
 The server uses Server-Sent Events (SSE) with content type `text/event-stream`.
 ### Anthropic-Compatible Endpoint
 The server also provides an Anthropic-compatible endpoint at `/v1/messages`:
 ```bash
 curl -X POST http://localhost:8000/v1/messages \
  -H "Content-Type: application/json" \
  -d '{
    "model": "astrai",
    "system": "You are a helpful assistant.",
    "messages": [{"role": "user", "content": "Hello, how are you?"}],
    "max_tokens": 2048
  }'
 ```
 Response:
 ```json
 {
  "id": "msg_abc123...",
  "type": "message",
  "role": "assistant",
  "model": "astrai",
  "content": [{"type": "text", "text": "Hello! I am doing well..."}],
  "stop_reason": "end_turn",
  "stop_sequence": null,
  "usage": {"input_tokens": 20, "output_tokens": 15}
 }
 ```
 Streaming:
 ```bash
 curl -X POST http://localhost:8000/v1/messages \
  -H "Content-Type: application/json" \
  -d '{
    "model": "astrai",
    "system": "You are a helpful assistant.",
    "messages": [{"role": "user", "content": "Write a short poem"}],
    "max_tokens": 500,
    "stream": true
  }'
 ```
 Supports `stop_sequences` for early termination:
 ```bash
 curl -X POST http://localhost:8000/v1/messages \
  -H "Content-Type: application/json" \
  -d '{
    "model": "astrai",
    "messages": [{"role": "user", "content": "Write a story"}],
    "max_tokens": 500,
    "stop_sequences": ["The end", "THE END"]
  }'
 ```
 ### Health Check
 Monitor server and model status:
 ```bash
 curl http://localhost:8000/health
 # {"status": "ok", "model_loaded": true}
 curl http://localhost:8000/stats
 # {"total_tasks": 10, "total_tokens": 5000, "active_tasks": 1, "waiting_queue": 0}
 ```
 > Document Update Time: 2026-04-09
@@ -1,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 |
+| `--max_grad_norm` | Maximum gradient norm for clipping (None disables) | None |
-### 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_len` | Maximum generation length | 1024 |
 | `stream` | Whether to stream output | False |
 ### Usage Example
 ```python
 import torch
 from astrai.model import AutoModel
 from astrai.tokenize import AutoTokenizer
 from astrai.inference import InferenceEngine, GenerationRequest
 # Load model using AutoModel
 model = AutoModel.from_pretrained("your_model_dir")
 # Load tokenizer
 tokenizer = AutoTokenizer.from_pretrained("your_model_dir")
 # Create engine with separate model and tokenizer
 engine = InferenceEngine(
    model=model,
    tokenizer=tokenizer,
 )
 # Build request with messages format
 request = GenerationRequest(
    messages=[
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": "Hello"},
    ],
    temperature=0.8,
    top_p=0.95,
    top_k=50,
    max_len=1024,
 )
 # Generate (streaming)
 for token in engine.generate_with_request(request):
    print(token, end="", flush=True)
 # Or use simple generate interface
 result = engine.generate(
    prompt="Hello",
    stream=False,
    max_tokens=1024,
    temperature=0.8,
    top_p=0.95,
    top_k=50,
 )
 ```
-### 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-04-09
+| Parameter | Type | Default | Description |
 |-----------|------|---------|-------------|
 | `--param_path` | str | required | Path to the model directory |
 | `--input_json_file` | str | required | Path to the input JSONL file |
 | `--output_json_file` | str | required | Path to the output JSONL file |
 | `--question_key` | str | `question` | Key for the question in input JSON |
 | `--response_key` | str | `response` | Key for the response in output JSON |
 | `--temperature` | float | `0.60` | Sampling temperature |
 | `--top_k` | int | `30` | Top-k filtering |
 | `--top_p` | float | `0.95` | Nucleus sampling threshold |
 | `--batch_size` | int | `1` | Batch size for generation |
 | `--max_tokens` | int | model config `max_len` | Maximum tokens to generate |
 Usage:
 ```bash
 python scripts/tools/generate.py \
    --param_path ./params \
    --input_json_file input.jsonl \
    --output_json_file output.jsonl
 ```
 ## Preprocess (`preprocess.py`)
 | Parameter | Type | Default | Description |
 |-----------|------|---------|-------------|
 | `input_files` | path(s) | required | Input JSONL file(s), supports glob (`data/*.jsonl`) |
 | `--output_dir`, `-o` | path | required | Output directory for processed data |
 | `--config`, `-c` | path | required | Preprocessing pipeline config (JSON) |
 | `--tokenizer_path` | str | `params` | Path to tokenizer directory |
 Usage:
 ```bash
 python scripts/tools/preprocess.py data/*.jsonl -o output/ -c sft.json
 ```
 See [Preprocessing Guide](preprocessing.md) for config file format and examples.
 ---
 > Document Update Time: 2026-07-19
@@ -0,0 +1,364 @@
 # Preprocessing Pipeline
 Declarative JSON-driven data preprocessing. `MaskBuilderFactory` supports three registered builders: `"single"` (single-output via `input.sections`), `"multi"` (multi-output via `input.sources`), and `"sectioned"` (façade dispatching to `single` or `multi` based on config).
 ## Contents
 - [Philosophy](#philosophy)
 - [Config Structure](#config-structure)
 - [Quick Start](#quick-start) — SFT Chat, SFT Instruction, Pretrain, DPO, GRPO examples
 - [Configuration Reference](#configuration-reference) — all fields
 - [Mask Algorithm](#mask-algorithm)
 - [Output Layout](#output-layout)
 - [CLI](#cli)
 - [Python API](#python-api)
 ## Philosophy
 | Component | Responsibility |
 |-----------|---------------|
 | `tokenizer_config.json` (`chat_template`) | Formatting -- how roles become tokens |
 | `pipeline.json` (`mask`) | Masking -- which roles participate in training |
 A single config file captures the entire pipeline, reusable and version-controllable.
 ## Config Structure
 ```json
 {
  "version":       1,
  "input":         {},   // sections (single) or sources (multi)
  "mask":          {},   // role -> "train" | "mask"
  "mask_default":  "mask",
  "preprocessing": {},
  "output":        {}
 }
 ```
 ### Section Fields
 | Field | Type | Default | Description |
 |-------|------|---------|-------------|
 | `field` | str | -- | JSONL key to read |
 | `action` | str | -- | `"train"` / `"mask"` / `"$role"` |
 | `template` | bool | `false` | Apply `chat_template` per message |
 | `add_special_tokens` | bool | `true` for first non-template section | Add special tokens during encode |
 ### Source Fields (multi-output mode)
 | Field | Type | Default | Description |
 |-------|------|---------|-------------|
 | `sections` | list[dict] | -- | Same as single-output section list |
 | `list_field` | bool | `false` | JSONL field holds a list; tokenise each element |
 | `mask_key` | str | `"{key}_mask"` | Explicit output key for loss mask |
 ---
 ## Quick Start
 ### SFT Chat
 Input JSONL:
 ```json
 {"messages": [{"role": "system", "content": "You are helpful."}, {"role": "user", "content": "Hi"}, {"role": "assistant", "content": "Hello!"}]}
 ```
 Config:
 ```json
 {
  "input": {
    "sections": [
      {"field": "messages", "action": "$role", "template": true}
    ]
  },
  "mask": {
    "system": "mask",
    "user": "mask",
    "assistant": "train"
  },
  "mask_default": "mask",
  "preprocessing": {
    "max_seq_len": 2048
  },
  "output": {
    "storage_format": "bin",
    "dtype": {"loss_mask": "bool"}
  }
 }
 ```
 Output keys: `sequence` (int32), `loss_mask` (bool)
 ### SFT Instruction
 Input JSONL:
 ```json
 {"prompt": "Translate to French: Hello", "response": "Bonjour"}
 ```
 Config:
 ```json
 {
  "input": {
    "sections": [
      {"field": "prompt",   "action": "mask", "add_special_tokens": true},
      {"field": "response", "action": "train"}
    ]
  },
  "mask_default": "mask",
  "preprocessing": {
    "max_seq_len": 2048
  }
 }
 ```
 Output keys: `sequence`, `loss_mask`
 ### Pretrain
 Input JSONL:
 ```json
 {"text": "Artificial Intelligence is a field of computer science..."}
 ```
 Config:
 ```json
 {
  "input": {
    "sections": [
      {"field": "text", "action": "train"}
    ]
  },
  "preprocessing": {
    "max_seq_len": 8192,
    "min_chars": 100
  }
 }
 ```
 Output keys: `sequence` (no `loss_mask` — all tokens trained)
 ### DPO
 Input JSONL:
 ```json
 {"chosen": [{"role": "user", "content": "What is 2+2?"}, {"role": "assistant", "content": "4"}], "rejected": [{"role": "user", "content": "What is 2+2?"}, {"role": "assistant", "content": "5"}]}
 ```
 Config:
 ```json
 {
  "input": {
    "sources": {
      "chosen": {
        "sections": [
          {"field": "chosen", "action": "$role", "template": true}
        ]
      },
      "rejected": {
        "sections": [
          {"field": "rejected", "action": "$role", "template": true}
        ]
      }
    }
  },
  "mask": {
    "user": "mask",
    "assistant": "train"
  },
  "mask_default": "mask"
 }
 ```
 Output keys: `chosen`, `chosen_mask`, `rejected`, `rejected_mask`
 ### GRPO
 Input JSONL:
 ```json
 {"prompt": [{"role": "user", "content": "What is 2+2?"}], "responses": ["4", "Five", "Four"], "rewards": [1.0, 0.3, 0.8]}
 ```
 Config:
 ```json
 {
  "input": {
    "sources": {
      "prompts": {
        "sections": [
          {"field": "prompt", "action": "mask", "template": true}
        ]
      },
      "responses": {
        "sections": [
          {"field": "responses", "action": "train"}
        ],
        "list_field": true,
        "mask_key": "masks"
      },
      "rewards": {
        "sections": [
          {"field": "rewards", "action": "value"}
        ]
      }
    }
  },
  "mask": {
    "user": "mask",
    "assistant": "train"
  },
  "mask_default": "mask"
 }
 ```
 Output keys: `prompts`, `prompts_mask`, `responses`, `masks`, `rewards` (float32)
 - `action: "value"` — extract raw values from JSONL without tokenisation
 - `list_field: true` — tokenise each list element independently, then concatenate
 - `mask_key: "masks"` — rename the auto-generated mask key (default: `responses_mask`)
 - `prompts_mask` is auto-generated (all masked) and unused by GRPOStrategy
 ---
 ## Configuration Reference
 ### `input`
 | Field | Type | Default | Description |
 |-------|------|---------|-------------|
 | `sections` | list[dict] or null | `null` | Section specs for single-output mode |
 | `sources` | dict[str, dict] or null | `null` | Source specs for multi-output mode (DPO/GRPO) |
 When `sources` is set, `sections` is ignored.
 ### `mask`
 | Field | Type | Default | Description |
 |-------|------|---------|-------------|
 | `mask` | dict | `{}` | `{role: "train" \| "mask"}` |
 | `mask_default` | str | `"mask"` | Default action for unlisted roles |
 ### `preprocessing`
 | Field | Type | Default | Description |
 |-------|------|---------|-------------|
 | `max_seq_len` | int | `2048` | Truncate sequences to this length |
 | `min_chars` | int | `50` | Skip text-mode items shorter than this |
 | `max_chars` | int | `2000000` | Skip text-mode items longer than this |
 | `max_items` | int or null | `null` | Stop after N documents |
 | `packing_strategy` | str | `"simple"` | Packing strategy: `"simple"`, `"bfd"`, `"bfd_split"` |
 | `max_packed_len` | int | `8192` | Maximum length of a packed bin |
 | `truncation_mode` | str | `"keep_start"` | How to truncate sequences: `"keep_start"` or `"keep_end"` |
 ### `output`
 | Field | Type | Default | Description |
 |-------|------|---------|-------------|
 | `domain_key` | str or null | `null` | JSONL key for domain grouping |
 | `storage_format` | str | `"bin"` | `"bin"` (mmap) or `"h5"` |
 | `max_tokens_per_shard` | int | `100000000` | Flush threshold in cumulative tokens |
 | `dtype` | dict[str, str] | `{}` | Per-key tensor dtype override (e.g. `{"loss_mask": "bool"}`) |
 | `position_ids_mode` | str | `"doc_reset"` | How to compute position_ids: `"none"`, `"doc_reset"`, `"continuous"` |
 ---
 ## Mask Algorithm
 ### Template mode (`template: true`)
 1. Prepend BOS token (masked)
 2. For each message in the field's array:
   1. Render through `chat_template` for that single message
   2. Encode rendered text
   3. Apply mask rule for the message's role
 ### Non-template mode
 Encode the field value as text. Mask value is 1 (train) or 0 (mask) per the section's `action`.
 ### Text config detection
 When no section uses `template` and all sections have `action: "train"`, the builder omits `loss_mask` from the output — all tokens are trained.
 ---
 ## Output Layout
 ### Single-Shard (`bin`)
 ```
 output/
  __default__/
    shard_0000/
      meta.json
      sequence.bin
      loss_mask.bin
  wiki/
    shard_0000/
      meta.json
      sequence.bin
      loss_mask.bin
 ```
 ### Multi-Shard (`bin`)
 When `max_tokens_per_shard` is exceeded:
 ```
 output/
  __default__/
    shard_0000/
      meta.json
      sequence.bin
      loss_mask.bin
    shard_0001/
      meta.json
      sequence.bin
      loss_mask.bin
 ```
 For `bin` format, `MmapStore` discovers all shards under the domain directory via `rglob("meta.json")`. For `h5` format, `H5Store` discovers `.h5`/`.hdf5` files via recursive glob.
 ---
 ## CLI
 ```bash
 # SFT
 python scripts/tools/preprocess.py data/sft/*.jsonl -o output/sft/ -c configs/sft_chat.json
 # DPO
 python scripts/tools/preprocess.py data/dpo/*.jsonl -o output/dpo/ -c configs/dpo.json --tokenizer_path params
 # GRPO
 python scripts/tools/preprocess.py data/grpo/*.jsonl -o output/grpo/ -c configs/grpo.json
 ```
 ---
 ## Python API
 ```python
 from astrai.preprocessing.pipeline import Pipeline
 from astrai.config.preprocess_config import PipelineConfig
 config = PipelineConfig.from_file("sft.json")
 Pipeline(
    config,
    ["data_part1.jsonl", "data_part2.jsonl"],
    output_dir="output/",
    tokenizer_path="params",
 ).run()
 ```
 > Document Update Time: 2026-07-09
@@ -0,0 +1,225 @@
 # Training
 ## Contents
 - [Autoregression](#autoregression)
 - [Causal Mask](#causal-mask)
 - [Rotary Position Embedding (RoPE)](#rotary-position-embedding-rope)
 - [Training Loop](#training-loop)
 - [Strategies](#strategies) — SEQ, SFT, DPO, GRPO
 - [LR Schedulers](#lr-schedulers)
 - [Gradient Checkpointing](#gradient-checkpointing)
 - [Checkpoint](#checkpoint)
 - [TrainContextBuilder](#traincontextbuilder-builder-pattern)
 - [Training CLI](#training-cli)
 ### Autoregression
 Given a token sequence, the model predicts the probability of the next token. Each generated token is appended to the input and fed back, repeating until an end-of-sequence token or max length.
 ### Causal Mask
 ```
 sequence : [[1, 2, 3, 4, 5, 6]]
 input_ids: [[1, 2, 3, 4, 5]]
 target_ids: [[2, 3, 4, 5, 6]]
 ```
 Lower-triangular mask prevents attending to future positions:
 ```
 [[0, -inf, -inf, -inf, -inf],
 [0,    0, -inf, -inf, -inf],
 [0,    0,    0, -inf, -inf],
 [0,    0,    0,    0, -inf],
 [0,    0,    0,    0,    0]]
 ```
 ### Rotary Position Embedding (RoPE)
 RoPE embeds position into Q/K vectors via complex rotation:
 $$ q_i = R_i W_q x_i, \quad k_j = R_j W_k x_j, \quad q_i^T k_j = x_i^T W_q^T R_{i-j} W_k x_j $$
 The complex rotation `freqs_cis` is pre-computed once (`cos, sin` pairs per position). `apply_rotary_emb` multiplies Q/K as complex numbers.
 ## Training Loop
 Two-level loop: **epoch** → **batch**. Optimizer step fires every `grad_accum_steps` batches.
 ```
 on_train_begin
  model.train()
  on_epoch_begin
    for batch in dataloader:
      on_batch_begin
      with executor.accumulate(model):
        loss = strategy.compute_loss(batch)
        context.loss = loss.item()
        stand_loss = loss / executor.grad_accum_steps
        executor.backward(stand_loss)
        context.consumed_samples += (
            context.config.batch_per_device * context.world_size
        )
        on_batch_end
        if executor.sync_gradients:
          on_optimizer_step
          optimizer.step()
          optimizer.zero_grad()
          if scheduler:
            scheduler.step()
    on_epoch_end
 on_train_end
 ```
 ### Callback Lifecycle
 | Hook | Fires | Default callback |
 |------|-------|-----------------|
 | `on_train_begin` | Before training starts | `GradientCheckpointingCallback` |
 | `on_epoch_begin` | Start of each epoch | `ProgressBarCallback` |
 | `on_batch_begin` | Every batch | — |
 | `on_optimizer_step` | Every accumulation window | `GradientClippingCallback`, `MetricCallback`, `ProgressBarCallback` |
 | `on_batch_end` | Every batch | `CheckpointCallback` |
 | `on_epoch_end` | End of each epoch | `MetricCallback`, `ProgressBarCallback` |
 | `on_error` | On exception during training | `CheckpointCallback`, `MetricCallback` |
 | `on_train_end` | Training ends (always via finally) | `CheckpointCallback`, `MetricCallback`, `GradientCheckpointingCallback` |
 Default callbacks (in order): `gradient_checkpointing` (activation checkpointing, optional), `checkpoint` (safetensors, rank-0), `metric` (JSONL + validation, rank-0), `progress_bar` (tqdm), `gradient_clipping` (always registered; computes grad norm, clips only when `max_grad_norm` is not `None`).
 ## Strategies
 ### SEQ (Pre-training)
 Next-token cross-entropy with optional label smoothing:
 $$
 L_{\text{PT}} = -\sum_{t=1}^{T} \log P(x_t \mid x_{\lt t}; \theta)
 $$
 Keys: `input_ids`, `target_ids`. Optional: `label_smoothing`.
 ### SFT (Supervised Fine-Tuning)
 Masked cross-entropy (`ignore_index=-100`) over response tokens:
 $$
 L_{\text{SFT}} = -\sum_{t=P+1}^{P+L} \log P(s_t \mid s_{\lt t}; \theta)
 $$
 Keys: `input_ids`, `target_ids`, `loss_mask`, `position_ids`. Optional: `label_smoothing`.
 ### DPO (Direct Preference Optimization)
 Frozen reference model, preference margin via log-ratio:
 $$
 L_{\text{DPO}} = -\mathbb{E}\left[\log\sigma\left(\beta\log\frac{\pi_\theta(y_w\mid x)}{\pi_{\text{ref}}(y_w\mid x)} - \beta\log\frac{\pi_\theta(y_l\mid x)}{\pi_{\text{ref}}(y_l\mid x)}\right)\right]
 $$
 Parameters: `beta=0.1`, `reduction="sum"`. Keys: `chosen`, `rejected`, `chosen_mask`, `rejected_mask`.
 ### GRPO (Group Relative Policy Optimization)
 Token-level PPO with group-normalized advantages. Advantages are derived from
 scalar per-response rewards, group-normalized, and broadcast across all response
 tokens. Only response tokens contribute to the loss (prompt tokens are masked
 out):
 $$
 \text{Advantage}_i = \frac{r_i - \mu}{\sigma + \epsilon}
 $$
 $$
 L_{\text{GRPO}} = -\mathbb{E}_t\left[\min\left(\rho_t A,\; \text{clip}\left(\rho_t, 1-\epsilon, 1+\epsilon\right)A\right)\right] + \lambda \cdot \mathbb{E}_t\left[\frac{\pi_{\text{ref}}}{\pi_\theta} - \log\frac{\pi_{\text{ref}}}{\pi_\theta} - 1\right]
 $$
 where $\rho_t = \pi_\theta(a_t|s_t) / \pi_{\text{old}}(a_t|s_t)$ is the
 per-token importance sampling ratio against the behaviour policy
 (`old_model`, synced externally between data-generation rounds) and the
 expectations are over valid response tokens. The KL term regularises
 $\pi_\theta$ towards a frozen reference model (`ref_model`, typically
 the SFT checkpoint).
 Parameters: `group_size=4`, `clip_eps=0.2`, `kl_coef=0.01`. External sync of `old_model` weights via `sync_old_model()` between data-generation rounds.
 Keys: `prompts`, `responses`, `masks`, `rewards`.
 ## LR Schedulers
 | Type | Class | Description |
 |------|-------|-------------|
 | Cosine | `CosineScheduler` | Linear warmup → cosine decay to `min_rate` |
 | SGDR | `SGDRScheduler` | Cosine annealing with warm restarts (`t_mult=2`) |
 | WSD | `WSDScheduler` | Warmup-Stable-Decay with sqrt cooldown |
 Created by `SchedulerFactory.create(schedule_type, optimizer, **kwargs)`. Valid types: `"cosine"`, `"sgdr"`, `"wsd"`. Omit to use no scheduler.
 ## Gradient Checkpointing
 Trades compute for memory by recomputing activations during backward pass. Specify module types via `gradient_checkpointing_modules`:
 ```python
 from astrai.model.components.decoder_block import DecoderBlock
 config = TrainConfig(..., gradient_checkpointing_modules=[DecoderBlock])
 ```
 Callback wraps each `DecoderBlock.forward` with `torch.utils.checkpoint.checkpoint(use_reentrant=False)`, compatible with `torch.compile`. Uses `nn.Module.apply()` for traversal — works through DDP wrappers without manual unwrap. Empty list (default) means no-op.
 ## Checkpoint
 ```
 Checkpoint(state_dict, epoch, consumed_samples, extra, meta, config)
  ├── save(save_dir)    rank-0 only: meta.json (epoch/consumed_samples/timestamp) + config.json (model config) + model.safetensors + optional {key}.pt (optimizer.pt, scheduler.pt)
  └── load(save_dir, broadcast=False)    loads from local disk; set broadcast=True to broadcast metadata from rank-0
 ```
 Optimizer/scheduler state persisted by default via `Checkpoint.extra`.  
 Model config (`context.model_config`) saved into `config.json` during training via `CheckpointCallback`.
 ## TrainContextBuilder (Builder Pattern)
 ```python
 context = (
    TrainContextBuilder(config)
        .with_resume_dir(resume_dir)
        .build()
 )
 # Returns TrainContext with model, strategy, optimizer, scheduler, dataloader, checkpoint
 ```
 - Loads checkpoint weights if provided
 - Creates executor via `ExecutorFactory.create(cfg.parallel_mode, grad_accum_steps=cfg.grad_accum_steps, **cfg.executor_kwargs)`
 - Calls `executor.prepare(model, optimizer, dataloader, scheduler)` for model distribution (e.g. DDP) + gradient accumulation wrappers
 - Creates `ResumableDistributedSampler` for shuffle+resume
 - Builds strategy via `StrategyFactory.create(train_type, model, device, **kwargs)`
 ## Training CLI
 ```bash
 export CUDA_VISIBLE_DEVICES=0,1,2,3
 nohup python scripts/tools/train.py \
    --nprocs=4 \
    --parallel_mode=ddp \
    --train_type=seq \
    --data_root_path=/path/to/dataset \
    --param_path=/path/to/model \
    --batch_per_device=4 \
    --grad_accum_steps=8 \
    --warmup_ratio=0.05 \
    --max_lr=1e-4 \
    --max_grad_norm=1.0 \
    --weight_decay=0.1 \
    --window_size=2048 \
    --ckpt_interval=10000 \
    --ckpt_dir=./checkpoint \
    --random_seed=3407 \
    --label_smoothing=0.05 \
    > out.log 2> err.log &
 ```
 Full parameter reference at [params.md](params.md).
 > Document Update Time: 2026-07-19
@@ -1,32 +1,98 @@
-__version__ = "1.3.4"
+__version__ = "1.3.10"
 __author__ = "ViperEkura"
 from astrai.config import (
-    ModelConfig,
+    AutoRegressiveLMConfig,
    BaseModelConfig,
    ConfigFactory,
    EncoderConfig,
    PipelineConfig,
    TrainConfig,
 )
-from astrai.dataset import DatasetFactory
+from astrai.dataset import (
    BaseDataset,
    DatasetFactory,
    RDSampler,
    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",
    "RDSampler",
    "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,80 @@
-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(
+    max_grad_norm: Optional[float] = field(
-        default=1.0, metadata={"help": "Maximum gradient norm."}
+        default=None,
        metadata={"help": "Maximum gradient norm. None disables clipping."},
    )
    gradient_checkpointing_modules: List[str] = field(
        default_factory=list,
        metadata={"help": "Module types to enable activation checkpointing for."},
    )
    # 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
@@ -51,6 +88,10 @@ class TrainConfig:
    pin_memory: bool = field(
        default=False, metadata={"help": "Pin memory for dataloader."}
    )
    collate_fn: Optional[Callable[[List[Any]], Any]] = field(
        default=None,
        metadata={"help": "Collate function for dataloader (e.g. dpo_collate_fn)."},
    )
    # distributed training
    nprocs: int = field(
@@ -66,18 +107,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 +150,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.")
@@ -1,19 +1,43 @@
 from astrai.dataset.dataset import (
    BaseDataset,
    BaseSegmentFetcher,
    DatasetFactory,
-    MultiSegmentFetcher,
+    dpo_collate_fn,
    grpo_collate_fn,
 )
 from astrai.dataset.sampler import RDSampler
 from astrai.dataset.storage import (
    H5Store,
    JsonlStore,
    MmapStore,
    Recordable,
    Store,
    StoreFactory,
    Streamable,
    detect_format,
 )
 from astrai.serialization import (
    load_bin,
    load_h5,
    save_bin,
    save_h5,
 )
 from astrai.dataset.sampler import ResumableDistributedSampler
 __all__ = [
    # Base classes
    "BaseDataset",
    # Factory
    "DatasetFactory",
-    # Fetchers
+    "dpo_collate_fn",
-    "BaseSegmentFetcher",
+    "grpo_collate_fn",
-    "MultiSegmentFetcher",
+    "Store",
-    # Sampler
+    "Streamable",
-    "ResumableDistributedSampler",
+    "Recordable",
    "StoreFactory",
    "H5Store",
    "MmapStore",
    "JsonlStore",
    "detect_format",
    "save_h5",
    "load_h5",
    "save_bin",
    "load_bin",
    "RDSampler",
 ]
@@ -1,308 +1,176 @@
-"""Dataset implementations with factory pattern for training."""
+"""Dataset implementations for training.
 Composition over inheritance — every dataset is a thin wrapper that
 binds a :class:`Store` to a particular train-type's key mapping.  All
 sample-id → token/record indexing lives on the Store; datasets never
 know about window/stride math or segment layouts.
 Class hierarchy:
    BaseDataset (ABC)         — holds a Store, exposes __len__/keys,
                                 overrides __getitem__
    ├── SEQDataset            — next-token prediction (stream)
    ├── SFTDataset            — loss-mask + position_ids (stream)
    ├── DPODataset            — chosen/rejected pairs (record)
    └── GRPODataset           — prompt + response group (record)
 ``DatasetFactory.load(train_type, load_path, window_size, stride, …)``
 builds the Store (auto-detecting format) before constructing the
 matching dataset.  Passing ``store=`` skips Store construction.
 When a record dataset (DPO) reads from raw JSONL, a *processor*
 function (pure ``record -> Dict[str, Tensor]``) is forwarded to
 :class:`JsonlStore` so tokenisation happens on the fly.
 """
 import bisect
 from abc import ABC, abstractmethod
-from typing import Dict, List, Optional, Union
+from functools import partial
 from typing import Callable, Dict, List, Optional
 import torch
 from torch import Tensor
 from torch.utils.data import Dataset
 from astrai.dataset.storage import (
    Store,
    StoreFactory,
    detect_format,
 )
 from astrai.factory import BaseFactory
-from astrai.serialization import load_h5
+from astrai.tokenize import AutoTokenizer
-class BaseSegmentFetcher:
+def dpo_tokenize(
-    """Fetches data segments across multiple tensor segments.
+    record: dict,
    tokenizer,
    max_len: int = 2048,
 ) -> Optional[dict]:
    """Tokenize one DPO record into chosen/rejected + masks.
-    Maintains cumulative lengths for efficient range queries across
+    Applies the tokenizer's chat template so token sequences match the
-    multiple discontinuous segments.
+    SFT checkpoint's format.  Prompt is rendered with
    ``add_generation_prompt=True``; chosen/rejected are appended as a
    single assistant turn.
    Accepts:
    - Flat:   ``{"prompt": str, "chosen": str, "rejected": str}``
    - Conv:   ``{"prompt": [{role, content}, ...], "chosen": [...], ...}``
    - Legacy: ``{"input": str, "chosen": str, "rejected": str}``
    No packing, no ``position_ids`` — DPO sequences are independent.
    """
    prompt = record.get("prompt") or record.get("input")
    chosen = record.get("chosen")
    rejected = record.get("rejected")
    if prompt is None or chosen is None or rejected is None:
        return None
-    def __init__(self, segments: List[Tensor]):
+    prompt_messages = _to_messages(prompt)
-        self.segments = segments
+    chosen_text = _extract_text(chosen)
-        self.cum_lengths = []
+    rejected_text = _extract_text(rejected)
    if chosen_text is None or rejected_text is None:
        return None
    chosen_messages = prompt_messages + [{"role": "assistant", "content": chosen_text}]
    rejected_messages = prompt_messages + [
        {"role": "assistant", "content": rejected_text}
    ]
-        total = 0
+    prompt_ids = tokenizer.apply_chat_template(
-        for seg in segments:
+        prompt_messages, tokenize=True, add_generation_prompt=True
-            total += torch.numel(seg)
+    )
-            self.cum_lengths.append(total)
+    ch_ids = tokenizer.apply_chat_template(
        chosen_messages, tokenize=True, add_generation_prompt=False
    )
    re_ids = tokenizer.apply_chat_template(
        rejected_messages, tokenize=True, add_generation_prompt=False
    )
-        self.total_length = total
+    full_ch = ch_ids[:max_len]
    full_re = re_ids[:max_len]
-    def __len__(self) -> int:
+    prompt_len = min(len(prompt_ids), max_len)
-        return self.total_length
+    ch_mask = [0] * prompt_len + [1] * max(0, len(full_ch) - prompt_len)
    ch_mask = ch_mask[:max_len]
    re_mask = [0] * prompt_len + [1] * max(0, len(full_re) - prompt_len)
    re_mask = re_mask[:max_len]
-    def fetch_data(self, begin_idx: int, end_idx: int) -> Tensor:
+    return {
-        """Fetch data in the range [begin_idx, end_idx).
+        "chosen": full_ch,
-
+        "rejected": full_re,
-        Args:
+        "chosen_mask": ch_mask,
-            begin_idx: Starting index (inclusive)
+        "rejected_mask": re_mask,
            end_idx: Ending index (exclusive)
        Returns:
            Concatenated tensor of data in the specified range
        """
        if not (
            0 <= begin_idx < self.total_length and 0 <= end_idx <= self.total_length
        ):
            raise ValueError("begin_idx or end_idx out of bounds")
        if begin_idx >= end_idx:
            return torch.tensor([], dtype=torch.long)
        # Find segment boundaries for the range
        seg_start_idx = bisect.bisect_right(self.cum_lengths, begin_idx)
        seg_end_idx = bisect.bisect_left(self.cum_lengths, end_idx)
        result_segments = []
        for i in range(seg_start_idx, seg_end_idx + 1):
            prev_cum = self.cum_lengths[i - 1] if i > 0 else 0
            start = max(begin_idx - prev_cum, 0)
            end = min(end_idx - prev_cum, len(self.segments[i]))
            data = self.segments[i][start:end]
            result_segments.append(data)
        return torch.cat(result_segments, dim=0)
 class MultiSegmentFetcher:
    """Manages multiple segment fetchers for different data keys.
    Each key corresponds to a different type of data (e.g., "sequence", "mask").
    """
    def __init__(self, multi_segments: Dict):
        self.multi_keys = list(multi_segments.keys())
        self.multi_fetchers = {
            key: BaseSegmentFetcher(segments)
            for key, segments in multi_segments.items()
    }
    def __len__(self) -> int:
        """Returns the minimum length across all fetchers."""
        len_list = [len(seg) for seg in self.multi_fetchers.values()]
        return min(len_list)
-    def key_fetch(
+def _to_messages(value) -> list:
-        self, begin_idx: int, end_idx: int, keys: Union[str, List[str]]
+    """Accept str or conversation list; return message list."""
-    ) -> Dict:
+    if isinstance(value, str):
-        """Fetch data for specific keys.
+        return [{"role": "user", "content": value}]
    if isinstance(value, list):
        return value
    return [{"role": "user", "content": str(value)}]
        Args:
            begin_idx: Starting index
            end_idx: Ending index
            keys: Single key or list of keys to fetch
-        Returns:
+def _extract_text(value) -> Optional[str]:
-            Dictionary of tensors if multiple keys, single tensor if one key
+    """Accept str or conversation list; return plain text."""
    if value is None:
        return None
    if isinstance(value, str):
        return value
    if isinstance(value, list):
        return "".join(m.get("content", "") for m in value if isinstance(m, dict))
    return None
 def dpo_processor(
    record: dict,
    tokenizer,
    max_len: int = 2048,
 ) -> Dict[str, Tensor]:
    """DPO processor: wraps :func:`dpo_tokenize` and returns tensors."""
    result = dpo_tokenize(record, tokenizer, max_len=max_len)
    if result is None:
        raise ValueError(f"Malformed DPO record: {list(record.keys())}")
    return {
        "chosen": torch.tensor(result["chosen"], dtype=torch.int32),
        "rejected": torch.tensor(result["rejected"], dtype=torch.int32),
        "chosen_mask": torch.tensor(result["chosen_mask"], dtype=torch.bool),
        "rejected_mask": torch.tensor(result["rejected_mask"], dtype=torch.bool),
    }
 def dpo_collate_fn(batch: List[Dict[str, Tensor]]) -> Dict[str, Tensor]:
    """Collate variable-length DPO samples into padded 2-D tensors.
    Input: list of dicts, each with:
      - chosen:        [C_i]
      - rejected:      [R_i]
      - chosen_mask:   [C_i]
      - rejected_mask: [R_i]
    Output (padded to the max length across chosen/rejected within the batch):
      - chosen:        [B, S_max]
      - rejected:      [B, S_max]
      - chosen_mask:   [B, S_max]
      - rejected_mask: [B, S_max]
    """
-        fetch_dict = {}
+    B = len(batch)
-        keys = [keys] if isinstance(keys, str) else keys
+    S_max = max(b["chosen"].size(0) for b in batch)
    S_max = max(S_max, max(b["rejected"].size(0) for b in batch))
-        for key in keys:
+    chosen = torch.zeros(B, S_max, dtype=torch.long)
-            fetcher = self.multi_fetchers[key]
+    rejected = torch.zeros(B, S_max, dtype=torch.long)
-            fetch_tensor = fetcher.fetch_data(begin_idx, end_idx)
+    chosen_mask = torch.zeros(B, S_max, dtype=torch.bool)
-            fetch_dict[key] = fetch_tensor
+    rejected_mask = torch.zeros(B, S_max, dtype=torch.bool)
-        return fetch_dict if len(keys) > 1 else fetch_dict[keys[0]]
+    for i, b in enumerate(batch):
-
+        c_len = b["chosen"].size(0)
-    def fetch_data(self, begin_idx: int, end_idx: int) -> Dict:
+        r_len = b["rejected"].size(0)
-        """Fetch all keys."""
+        chosen[i, :c_len] = b["chosen"]
-        return self.key_fetch(begin_idx, end_idx, self.multi_keys)
+        rejected[i, :r_len] = b["rejected"]
-
+        chosen_mask[i, :c_len] = b["chosen_mask"]
-
+        rejected_mask[i, :r_len] = b["rejected_mask"]
 class BaseDataset(Dataset, ABC):
    """Abstract base class for all dataset types.
    Implements common functionality for window-based data fetching.
    """
    def __init__(self, window_size: int, stride: int):
        super().__init__()
        self.segments = {}
        self.window_size = window_size
        self.stride = stride
        self.total_samples = None
        self.fetcher: Optional[MultiSegmentFetcher] = None
    def load(self, load_path: str):
        """Load dataset from HDF5 file.
        Args:
            load_path: Path to the HDF5 data file
        """
        self.segments = load_h5(load_path)
        self.fetcher = MultiSegmentFetcher(self.segments)
        self.total_samples = len(self.fetcher)
    def get_index(self, index: int) -> tuple:
        """Calculate begin and end indices for a sample.
        Args:
            index: Sample index
        Returns:
            Tuple of (begin_idx, end_idx)
        """
        assert self.total_samples > self.window_size
        begin_idx = min(index * self.stride, self.total_samples - 1 - self.window_size)
        end_idx = min(begin_idx + self.window_size, self.total_samples - 1)
        return begin_idx, end_idx
    @abstractmethod
    def __getitem__(self, index: int) -> Dict[str, Tensor]:
        """Get a single sample by index.
        Must be implemented by subclasses.
        """
        raise NotImplementedError
    def __len__(self) -> int:
        assert self.total_samples is not None
        if self.total_samples <= self.window_size:
            return 0
        return (self.total_samples - 1 - self.window_size) // self.stride + 1
 class DatasetFactory(BaseFactory["BaseDataset"]):
    """Factory class for creating dataset instances.
    Supports decorator-based registration for extensible dataset types.
    All default dataset types (seq, sft, dpo, grpo) are registered automatically
    when their classes are defined with the decorator.
    Example usage:
        @DatasetFactory.register("custom")
        class CustomDataset(BaseDataset):
            ...
        dataset = DatasetFactory.create("custom", window_size, stride)
    """
    @classmethod
    def _validate_component(cls, dataset_cls: type) -> None:
        """Validate that the dataset class inherits from BaseDataset."""
        if not issubclass(dataset_cls, BaseDataset):
            raise TypeError(f"{dataset_cls.__name__} must inherit from BaseDataset")
    @classmethod
    def create(cls, train_type: str, window_size: int, stride: int) -> "BaseDataset":
        """Create a dataset instance.
        Args:
            train_type: Type of training ("seq", "sft", "dpo", "grpo")
            window_size: Window size for data sampling
            stride: Stride between consecutive samples
        Returns:
            Dataset instance
        """
        return super().create(train_type, window_size, stride)
    @classmethod
    def load(
        cls,
        train_type: str,
        load_path: str,
        window_size: int,
        stride: Optional[int] = None,
    ) -> "BaseDataset":
        """Create and load a dataset in one step.
        Args:
            train_type: Type of training dataset
            load_path: Path to the data file
            window_size: Window size for data sampling
            stride: Stride between consecutive samples (default: same as window_size)
        Returns:
            Loaded dataset instance
        """
        if stride is None:
            stride = window_size
        dataset = cls.create(train_type, window_size, stride)
        dataset.load(load_path)
        return dataset
    @classmethod
    def available_types(cls) -> list:
        """Return list of registered dataset type names."""
        return cls.list_registered()
 # ============== Dataset Classes ==============
 # All dataset classes are registered at class definition time using the decorator
@DatasetFactory.register("seq")
 class SEQDataset(BaseDataset):
    """Dataset for sequential next-token prediction training."""
    def __init__(self, window_size: int, stride: int):
        super().__init__(window_size, stride)
    def _fetch_data(self, begin_idx: int, end_idx: int) -> Tensor:
        return self.fetcher.key_fetch(begin_idx, end_idx, "sequence")
    def __getitem__(self, index):
        begin_idx, end_idx = self.get_index(index)
        x = self._fetch_data(begin_idx, end_idx).to(dtype=torch.long)
        y = self._fetch_data(begin_idx + 1, end_idx + 1).to(dtype=torch.long)
        return {"input_ids": x, "target_ids": y}
@DatasetFactory.register("sft")
 class SFTDataset(BaseDataset):
    """Dataset for supervised fine-tuning with loss masking."""
    def __init__(self, window_size: int, stride: int):
        super().__init__(window_size, stride)
    def _fetch_data(self, begin_idx: int, end_idx: int, key: str) -> Tensor:
        return self.fetcher.key_fetch(begin_idx, end_idx, key)
    def __getitem__(self, index):
        begin_idx, end_idx = self.get_index(index)
        x = self._fetch_data(begin_idx, end_idx, "sequence").to(dtype=torch.long)
        y = self._fetch_data(begin_idx + 1, end_idx + 1, "sequence").to(
            dtype=torch.long
        )
        loss_mask = self._fetch_data(begin_idx + 1, end_idx + 1, "loss_mask").to(
            dtype=torch.bool
        )
        return {"input_ids": x, "target_ids": y, "loss_mask": loss_mask}
@DatasetFactory.register("dpo")
 class DPODataset(BaseDataset):
    """Dataset for Direct Preference Optimization training."""
    def __init__(self, window_size: int, stride: int):
        super().__init__(window_size, stride)
    def _fetch_data(self, begin_idx: int, end_idx: int, key: str) -> Tensor:
        return self.fetcher.key_fetch(begin_idx, end_idx, key)
    def __getitem__(self, index: int):
        begin_idx, end_idx = self.get_index(index)
        chosen = self._fetch_data(begin_idx, end_idx, "chosen").to(dtype=torch.long)
        rejected = self._fetch_data(begin_idx, end_idx, "rejected").to(dtype=torch.long)
        chosen_mask = self._fetch_data(begin_idx, end_idx, "chosen_mask").to(
            dtype=torch.bool
        )
        rejected_mask = self._fetch_data(begin_idx, end_idx, "rejected_mask").to(
            dtype=torch.bool
        )
    return {
        "chosen": chosen,
@@ -312,23 +180,40 @@ class DPODataset(BaseDataset):
    }
-@DatasetFactory.register("grpo")
+def grpo_collate_fn(batch: List[Dict[str, Tensor]]) -> Dict[str, Tensor]:
-class GRPODataset(BaseDataset):
+    """Collate variable-length GRPO samples into padded 3-D tensors.
    """Dataset for Group Relative Policy Optimization training."""
-    def __init__(self, window_size: int, stride: int):
+    Input: list of dicts, each with:
-        super().__init__(window_size, stride)
+      - prompts:  [P_i]
      - responses: list of G tensors, each [R_ij]
      - masks:     list of G tensors, each [R_ij]
      - rewards:  [G]
-    def _fetch_data(self, begin_idx: int, end_idx: int, key: str) -> Tensor:
+    Output:
-        return self.fetcher.key_fetch(begin_idx, end_idx, key)
+      - prompts:   [B, P_max]
      - responses: [B, G, R_max]
      - masks:     [B, G, R_max]
      - rewards:   [B, G]
    """
    B = len(batch)
    G = len(batch[0]["responses"])
    P_max = max(b["prompts"].size(0) for b in batch)
    R_max = max(r.size(0) for b in batch for r in b["responses"])
-    def __getitem__(self, index: int) -> Dict[str, Tensor]:
+    prompts = torch.zeros(B, P_max, dtype=torch.long)
-        begin_idx, end_idx = self.get_index(index)
+    responses = torch.zeros(B, G, R_max, dtype=torch.long)
    masks = torch.zeros(B, G, R_max, dtype=torch.bool)
    rewards = torch.zeros(B, G, dtype=torch.float32)
-        prompts = self._fetch_data(begin_idx, end_idx, "prompts")
+    for i, b in enumerate(batch):
-        responses = self._fetch_data(begin_idx, end_idx, "responses")
+        p_len = b["prompts"].size(0)
-        masks = self._fetch_data(begin_idx, end_idx, "masks")
+        prompts[i, :p_len] = b["prompts"]
-        rewards = self._fetch_data(begin_idx, end_idx, "rewards")
+        rewards[i, : b["rewards"].size(0)] = b["rewards"]
        for g in range(min(G, len(b["responses"]))):
            r_len = b["responses"][g].size(0)
            responses[i, g, :r_len] = b["responses"][g]
            if g < len(b["masks"]):
                masks[i, g, :r_len] = b["masks"][g]
    return {
        "prompts": prompts,
@@ -336,3 +221,274 @@ class GRPODataset(BaseDataset):
        "masks": masks,
        "rewards": rewards,
    }
 def validate_keys(store: Store, required: List[str]) -> None:
    """Raise ``KeyError`` if *store* is missing any *required* key."""
    if not required:
        return
    actual = set(store.keys)
    missing = [k for k in required if k not in actual]
    if missing:
        raise KeyError(
            f"Store at {getattr(store, '_load_path', '?')} is missing required "
            f"keys {missing}; available keys are {sorted(actual)}."
        )
 class BaseDataset(Dataset, ABC):
    """Abstract base class for dataset types.
    Holds a :class:`Store`.  All sample-id indexing is delegated to the
    store — this class exposes ``__len__`` as ``len(store)`` and the
    ``keys`` property as ``store.keys``.  Subclasses implement
    ``__getitem__`` with the train-type-specific key mapping and any
    training-only index arithmetic (e.g. the next-token ``+1`` shift).
    """
    required_keys: List[str] = []
    def __init__(self, store: Store):
        super().__init__()
        self.store: Store = store
        validate_keys(store, self.required_keys)
    def __len__(self) -> int:
        return len(self.store)
    @property
    def keys(self) -> List[str]:
        return self.store.keys
    @property
    def token_count(self) -> int:
        return self.store.token_count
    @abstractmethod
    def __getitem__(self, index: int) -> Dict[str, Tensor]:
        raise NotImplementedError
 class DatasetFactory(BaseFactory["BaseDataset"]):
    """Factory for creating dataset instances by train-type.
    Use :meth:`DatasetFactory.register("custom")` to register new
    dataset classes; they must inherit from :class:`BaseDataset`.
    """
    @classmethod
    def load(
        cls,
        train_type: str,
        load_path: Optional[str] = None,
        window_size: int = 0,
        stride: Optional[int] = None,
        storage_type: Optional[str] = None,
        tokenizer_path: Optional[str] = None,
        max_len: int = 2048,
        store: Optional[Store] = None,
        **kwargs,
    ) -> "BaseDataset":
        """Create and load a dataset in one step.
        Two entry points:
        - **store given**: bind it directly — the caller fully controls
          Store construction and processor setup.  *load_path*,
          *storage_type*, *tokenizer_path*, *window_size*, *stride* are
          ignored.
        - **store is None**: build a Store from *load_path*, auto-detecting
          format and constructing a processor when *tokenizer_path* is
          given for a record dataset on JSONL.
        Args:
            train_type: Registered dataset name ("seq", "sft", "dpo",
                "grpo", …).
            load_path: Path to the data file or directory (ignored if
                *store* is given).
            window_size: Stream window length — only meaningful for
                stream datasets (SEQ/SFT).  Record datasets ignore it.
            stride: Stride between consecutive stream samples
                (default: same as *window_size*).
            storage_type: Storage backend ("h5", "bin", "jsonl") or
                None for auto-detection.
            tokenizer_path: Path to tokenizer for lazy JSONL
                tokenisation (record datasets only).
            max_len: Max sequence length forwarded to processors.
            store: Pre-built, already-loaded Store instance.
            **kwargs: Extra arguments forwarded to ``store.load()``.
        Returns:
            Loaded dataset instance.
        """
        if store is not None:
            return cls.create(train_type, store=store)
        if load_path is None:
            raise ValueError("Either load_path or store must be provided")
        if storage_type is None:
            storage_type = detect_format(load_path)
        if stride is None:
            stride = window_size
        processor = cls._maybe_build_processor(
            train_type, storage_type, tokenizer_path, max_len
        )
        store_window = cls._store_window_for(train_type, window_size)
        store = StoreFactory.create(
            storage_type,
            window_size=store_window,
            stride=stride if stride else store_window,
        )
        if processor is not None:
            store.load(load_path, processor=processor, **kwargs)
        else:
            store.load(load_path, **kwargs)
        return cls.create(train_type, store=store)
    @staticmethod
    def _store_window_for(train_type: str, window_size: int) -> int:
        """Stream datasets consume ``window_size``; record datasets ignore it.
        Record datasets (dpo/grpo) treat each record as an independent
        training unit and never window, so the store is built with
        ``window_size=0`` and ``len(store)`` returns the record count.
        """
        if train_type in ("seq", "sft"):
            return window_size
        return 0
    @staticmethod
    def _maybe_build_processor(
        train_type: str,
        storage_type: str,
        tokenizer_path: Optional[str],
        max_len: int,
    ) -> Optional[Callable[[dict], Dict[str, Tensor]]]:
        """Build an on-the-fly tokenisation processor if applicable.
        Only raw JSONL + record datasets (DPO/GRPO) need a processor;
        pre-tokenised backends (H5/bin) and stream datasets (SEQ/SFT)
        return ``None`` so no tokenizer is loaded.
        """
        if tokenizer_path is None or storage_type != "jsonl":
            return None
        if train_type == "dpo":
            tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
            return partial(dpo_processor, tokenizer=tokenizer, max_len=max_len)
        return None
@DatasetFactory.register("seq")
 class SEQDataset(BaseDataset):
    """Dataset for sequential next-token prediction training.
    Stream mode: ``store.fetch(begin, end, "sequence")`` returns the
    input window; the +1 shifted call returns the next-token target.
    """
    required_keys = ["sequence"]
    def __getitem__(self, index: int):
        begin, end = self.store.sample_window(index)
        x = self.store.fetch(begin, end, "sequence")
        y = self.store.fetch(begin + 1, end + 1, "sequence")
        return {
            "input_ids": x.to(dtype=torch.long),
            "target_ids": y.to(dtype=torch.long),
        }
@DatasetFactory.register("sft")
 class SFTDataset(BaseDataset):
    """Dataset for supervised fine-tuning with loss masking.
    Stream mode: ``sequence``/``loss_mask``/``position_ids`` are sliced
    to the window.  ``loss_mask`` and ``target_ids`` use the +1 shifted
    slice so they align with the predicted positions.
    """
    required_keys = ["sequence", "loss_mask", "position_ids"]
    def __getitem__(self, index: int):
        begin, end = self.store.sample_window(index)
        x = self.store.fetch(begin, end, "sequence")
        y = self.store.fetch(begin + 1, end + 1, "sequence")
        position_ids = self.store.fetch(begin, end, "position_ids")
        loss_mask = self.store.fetch(begin + 1, end + 1, "loss_mask")
        return {
            "input_ids": x.to(dtype=torch.long),
            "target_ids": y.to(dtype=torch.long),
            "position_ids": position_ids.to(dtype=torch.long),
            "loss_mask": loss_mask.to(dtype=torch.bool),
        }
@DatasetFactory.register("dpo")
 class DPODataset(BaseDataset):
    """Record-structured dataset for Direct Preference Optimization.
    Each sample is one preference pair (chosen + rejected) and is an
    independent training unit — no windowing, stride, or cross-record
    concatenation.  This keeps each sequence self-contained so attention
    never leaks across preference pairs.
    Two loading paths (handled by :class:`DatasetFactory`):
    - **Pre-tokenized** (H5/bin): ``store.load(path)`` reads per-record
      tensors; ``__getitem__`` returns them directly.
    - **Raw JSONL** (``tokenizer_path=...``): builds a lazy processor
      via :func:`dpo_processor` that tokenises on the fly — no packing,
      no ``position_ids``.
    """
    required_keys = ["chosen", "rejected", "chosen_mask", "rejected_mask"]
    def make_processor(self, tokenizer, max_len: int):
        return partial(dpo_processor, tokenizer=tokenizer, max_len=max_len)
    def __getitem__(self, index: int) -> Dict[str, Tensor]:
        return {
            "chosen": self.store.fetch_record(index, "chosen").to(dtype=torch.long),
            "rejected": self.store.fetch_record(index, "rejected").to(dtype=torch.long),
            "chosen_mask": self.store.fetch_record(index, "chosen_mask").to(
                dtype=torch.bool
            ),
            "rejected_mask": self.store.fetch_record(index, "rejected_mask").to(
                dtype=torch.bool
            ),
        }
@DatasetFactory.register("grpo")
 class GRPODataset(BaseDataset):
    """Dataset for offline Group Relative Policy Optimization.
    Each sample is one prompt with its group of responses and scalar
    rewards — an independent training unit with no windowing or stride.
    Expected storage layout (produced by JsonlStore or pre-tokenized):
    - ``prompts``:   List[Tensor]  — one 1-D token tensor per record
    - ``responses``: List[List[Tensor]] — G response tensors per record
    - ``masks``:     List[List[Tensor]] — G mask tensors per record
    - ``rewards``:   List[Tensor]  — one 1-D float tensor (len G) per record
    """
    required_keys = ["prompts", "responses", "masks", "rewards"]
    def __getitem__(self, index: int) -> Dict[str, Tensor]:
        prompts = self.store.fetch_record(index, "prompts")
        responses = self.store.fetch_record(index, "responses")
        masks = self.store.fetch_record(index, "masks")
        rewards = self.store.fetch_record(index, "rewards")
        return {
            "prompts": prompts.to(dtype=torch.long),
            "responses": [r.to(dtype=torch.long) for r in responses],
            "masks": [m.to(dtype=torch.bool) for m in masks],
            "rewards": rewards.to(dtype=torch.float32),
        }
@@ -5,7 +5,15 @@ import torch.distributed as dist
 from torch.utils.data import Dataset, Sampler
-class ResumableDistributedSampler(Sampler[int]):
+class RDSampler(Sampler[int]):
    """Resumable Distributed Sampler.
    A distributed sampler that supports checkpoint-based resume: iteration
    state (epoch, position) is tracked so training can continue from the
    exact sample after a restart.  Shards the dataset across
    ``dist.world_size`` replicas with optional shuffling.
    """
    def __init__(
        self,
        data_source: Dataset,
@@ -43,6 +51,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 +82,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
@@ -0,0 +1,636 @@
 """Storage backends for different data formats.
 Architecture (composition over inheritance):
    Store (ABC)               — owns _data/_cum/_offsets bookkeeping
                                  + window_size/stride for sample-id
                                  indexing.  __getitem__/__len__ produce
                                  the smallest iterable unit so Dataset
                                  classes are pure delegators.
    Streamable (mixin)        — raw token slice  fetch(begin, end, keys)
    Recordable (mixin)        — raw record slice fetch_record(idx, keys)
    H5Store(Store, Streamable, Recordable)
    MmapStore(Store, Streamable, Recordable)
    JsonlStore(Store, Streamable, Recordable)
 Each mixin is a stateless trait that relies on ``self._data`` etc.
 provided by :class:`Store`.  Concrete stores mix in whichever access
 primitives they support — ``Store`` is the sole base class, so there is
 no diamond inheritance or MRO ambiguity.
 Sample-id indexing lives on :class:`Store`, not on the dataset:
 - **Stream mode** (``window_size > 0``):  ``len(store)`` returns the number
  of ``(window_size, stride)`` windows that fit in the token river;
  ``store[i]`` returns the *i*-th window as a dict of per-key tensors;
  ``store.sample_window(i)`` exposes the underlying ``(begin, end)``
  token slice for callers (e.g. next-token trainers) that need a +1
  shifted companion window.
 - **Record mode** (``num_records > 0``):  ``len(store)`` returns the
  record count; ``store[i]`` returns the *i*-th record dict.
 Raw token/record access via :meth:`fetch` / :meth:`fetch_record`
 remains available for low-level callers that want explicit index
 control.  ``store.token_count`` is the total stream token count (what
 ``len(store)`` used to mean in the legacy stream-only API).
 ``segments_are_records`` (class attribute on each Store subclass)
 tells ``_normalize`` whether segments are inherently per-record (H5/
 JSONL) or opaque shards (bin).  Record access for bin relies on
 ``_offsets`` instead.
 :class:`JsonlStore` supports a lazy mode (``processor=fn``) that keeps
 raw records and defers tokenisation to ``fetch_record`` — used by DPO
 to train directly from a ``.jsonl`` file without a pre-tokenised copy.
 """
 import bisect
 import glob
 import json
 import logging
 from abc import ABC, abstractmethod
 from pathlib import Path
 from typing import Callable, Dict, List, Optional, Tuple, Union
 import torch
 from torch import Tensor
 from astrai.factory import BaseFactory
 from astrai.preprocessing.transform import TokenizeTransform
 from astrai.serialization import (
    load_bin,
    load_bin_offsets,
    load_h5,
 )
 logger = logging.getLogger(__name__)
 def detect_format(load_path: str) -> str:
    """Auto-detect storage format from files in the directory.
    Args:
        load_path: Directory or file path
    Returns:
        Format string ("h5", "bin", "jsonl", or "processed")
    Raises:
        FileNotFoundError: If no supported data files are found
    """
    root = Path(load_path)
    if root.is_file():
        suffix = root.suffix.lower()
        if suffix in (".h5", ".hdf5"):
            return "h5"
        if suffix == ".jsonl":
            return "jsonl"
        raise ValueError(f"Unsupported file format: {suffix}")
    h5_files = [
        Path(p)
        for pattern in ("*.h5", "*.hdf5")
        for p in glob.glob(str(root / "**" / pattern), recursive=True)
    ]
    if h5_files:
        return "h5"
    bin_files = [Path(p) for p in glob.glob(str(root / "**" / "*.bin"), recursive=True)]
    if bin_files:
        has_meta = (root / "meta.json").exists() or len(
            [Path(p) for p in glob.glob(str(root / "**" / "meta.json"), recursive=True)]
        ) > 0
        if has_meta:
            return "bin"
    jsonl_files = [
        Path(p) for p in glob.glob(str(root / "**" / "*.jsonl"), recursive=True)
    ]
    if jsonl_files:
        return "jsonl"
    raise FileNotFoundError(f"No supported data files found at {load_path}")
 class Store(ABC):
    """Common base for all storage backends.
    A Store owns both its data layout AND its sample-id → token/record
    index translation.  Datasets are thin wrappers that bind a Store
    to a particular train-type's key mapping; they never know about
    window/stride math.
    Two iteration modes:
    - **Stream** (``window_size > 0``): data is treated as one long
      token river.  ``len(store)`` returns the number of windows;
      ``store[i]`` slices every stream-compatible key to window ``i``;
      ``store.sample_window(i)`` returns the ``(begin, end)`` token
      slice for callers needing a +1 shifted companion window.
    - **Record** (``num_records > 0``): data is per-record.
      ``len(store)`` returns ``num_records``; ``store[i]`` returns
      the *i*-th record as a dict.
    Raw token slicing is still available via :meth:`fetch` (mixed in
    by :class:`Streamable`) when a store has stream support configured.
    Raw record slicing via :meth:`fetch_record` (mixed in by
    :class:`Recordable`) when a store has record support.
    ``token_count`` exposes the raw total stream length — this is what
    ``len(store)`` returned in the legacy stream-only API and what
    stream-bound ``fetch`` uses for its bounds check.
    """
    segments_are_records: bool = False
    def __init__(
        self,
        window_size: int = 0,
        stride: Optional[int] = None,
    ):
        self._data: Dict[str, List[Tensor]] = {}
        self._cum: Dict[str, List[int]] = {}
        self._offsets: Dict[str, List[int]] = {}
        self._length: int = 0
        self._num_records: int = 0
        self._window_size: int = int(window_size)
        self._stride: int = int(stride) if stride is not None else int(window_size)
    @abstractmethod
    def load(self, path: str, **kwargs) -> None:
        raise NotImplementedError
    @property
    def keys(self) -> List[str]:
        return list(self._data.keys())
    @property
    def window_size(self) -> int:
        return self._window_size
    @property
    def stride(self) -> int:
        return self._stride
    @property
    def token_count(self) -> int:
        """Total tokens across all stream segments.
        Useful for the bounds-checked raw :meth:`fetch` and as the
        legacy ``len(store)`` value.
        """
        return self._length
    @property
    def num_records(self) -> int:
        """Number of records available via :meth:`fetch_record`.
        Non-zero only when the backing layout provides per-record
        indexing (H5/JSONL segments or bin ``_offsets``).
        """
        return self._num_records
    @property
    def num_samples(self) -> int:
        """Number of items produced by ``__getitem__``.
        Stream-mode wins when ``window_size > 0`` and there are tokens
        to slice; otherwise falls back to ``num_records``.
        """
        if self._window_size > 0 and self._length > 0:
            total = self._length
            w = self._window_size
            if total <= w:
                return 0
            return (total - 1 - w) // self._stride + 1
        return self._num_records
    def __len__(self) -> int:
        return self.num_samples
    def __getitem__(self, index: int) -> Dict[str, Tensor]:
        if index < 0:
            index += self.num_samples
        if not 0 <= index < self.num_samples:
            raise IndexError(
                f"Store index out of range: {index}, num_samples={self.num_samples}"
            )
        if self._window_size > 0 and self._length > 0:
            begin, end = self.sample_window(index)
            keys = self._stream_keys()
            return {k: self.fetch(begin, end, k) for k in keys}
        return self.fetch_record(index, self._record_keys())
    def sample_window(self, index: int) -> Tuple[int, int]:
        """Return ``(begin, end)`` token positions for stream sample *index*.
        The clipped tail keeps the last reachable window inside the
        token river instead of overshooting.  Caller is responsible
        for staying within :attr:`num_samples`: an out-of-range index
        raises ``IndexError``.
        """
        if self._window_size <= 0:
            raise RuntimeError("sample_window() requires window_size > 0 (stream mode)")
        if self._window_size <= 0 or self._length <= self._window_size:
            raise IndexError(
                f"Data too short for window: token_count={self._length}, "
                f"window_size={self._window_size}"
            )
        if not 0 <= index < self.num_samples:
            raise IndexError(
                f"Sample index out of range: {index}, num_samples={self.num_samples}"
            )
        total = self._length
        begin = min(index * self._stride, total - 1 - self._window_size)
        end = min(begin + self._window_size, total - 1)
        return begin, end
    def _stream_keys(self) -> List[str]:
        out: List[str] = []
        for k, tensors in self._data.items():
            if tensors and isinstance(tensors[0], list):
                continue
            out.append(k)
        return out
    def _record_keys(self) -> List[str]:
        return list(self._data.keys())
    def _normalize(
        self,
        raw: Dict[str, list],
        offsets: Optional[Dict[str, List[int]]] = None,
    ):
        """Register segments and pre-compute indices for both access modes.
        Stream mode: ``_cum[key]`` accumulates per-segment lengths so
        ``Streamable._fetch_stream_key`` can bisect across segments
        without concatenation.
        Record mode: if *offsets* is provided (bin layout),
        ``_offsets[key]`` stores cumulative per-record offsets into the
        single concatenated segment.  Otherwise, when
        ``segments_are_records`` is True (H5/JSONL), ``_data[key]`` is
        a per-record list and ``fetch_record`` indexes it directly.
        Nested keys (GRPO ``responses``/``masks`` as
        ``List[List[Tensor]]``) are stored as-is and excluded from both
        cumulative bookkeepings — they are only accessed record-by-record.
        """
        flat_lengths = []
        for key, tensors in raw.items():
            self._data[key] = tensors
            if not tensors:
                self._cum[key] = []
                flat_lengths.append(0)
                continue
            if isinstance(tensors[0], list):
                self._cum[key] = []
                continue
            cum = []
            total = 0
            for t in tensors:
                total += t.shape[0]
                cum.append(total)
            self._cum[key] = cum
            flat_lengths.append(cum[-1] if cum else 0)
        self._length = min(flat_lengths) if flat_lengths else 0
        valid_offsets: Dict[str, List[int]] = {}
        if offsets:
            for key, off in offsets.items():
                segs = self._data.get(key, [])
                if len(segs) == 1 and len(off) > 1:
                    valid_offsets[key] = off
                elif len(segs) > 1:
                    logger.warning(
                        "Key '%s' has %d segments with offsets — record mode "
                        "disabled for this key (multi-shard bin+offsets not "
                        "supported). Merge shards or use H5/JSONL.",
                        key,
                        len(segs),
                    )
        self._offsets = valid_offsets
        if valid_offsets:
            record_counts = [len(v) - 1 for v in valid_offsets.values()]
            self._num_records = min(record_counts) if record_counts else 0
        elif self.segments_are_records:
            per_record_counts = []
            for key, tensors in self._data.items():
                if tensors and isinstance(tensors[0], list):
                    continue
                per_record_counts.append(len(tensors))
            self._num_records = min(per_record_counts) if per_record_counts else 0
        else:
            self._num_records = 0
 class Streamable:
    """Mixin granting raw token-stream access via :meth:`fetch`.
    Stateless trait relying on ``self._data``, ``self._cum``,
    ``self._length`` maintained by :class:`Store`.  Stream mode is
    active when the owning store has ``window_size > 0``; for stores
    that can also serve record access (H5/JSONL/bin+offsets), the
    ``fetch_record`` API from :class:`Recordable` is used instead.
    """
    def fetch(
        self,
        begin: int,
        end: int,
        keys: Union[str, List[str]],
    ):
        return _stream_fetch(self, begin, end, keys)
 def _stream_fetch(self, begin: int, end: int, keys: Union[str, List[str]]):
    if not getattr(self, "_data", None):
        raise RuntimeError("Store not loaded")
    if not (0 <= begin < self._length and 0 <= end <= self._length):
        raise ValueError(
            f"Index out of bounds: begin={begin}, end={end}, length={self._length}"
        )
    if isinstance(keys, str):
        return _fetch_stream_key(self, keys, begin, end)
    return {k: _fetch_stream_key(self, k, begin, end) for k in keys}
 def _fetch_stream_key(self, key: str, begin: int, end: int) -> Tensor:
    segments = self._data[key]
    cum = self._cum[key]
    seg_start = bisect.bisect_right(cum, begin)
    seg_end = bisect.bisect_left(cum, end)
    results = []
    for i in range(seg_start, seg_end + 1):
        prev = cum[i - 1] if i > 0 else 0
        s = max(begin - prev, 0)
        e = min(end - prev, segments[i].shape[0])
        results.append(segments[i][s:e])
    return results[0] if len(results) == 1 else torch.cat(results, dim=0)
 class Recordable:
    """Mixin granting raw record access via :meth:`fetch_record`.
    Stateless trait relying on ``self._data``, ``self._offsets``,
    ``self._num_records`` maintained by :class:`Store`.
    """
    def fetch_record(
        self,
        index: int,
        keys: Union[str, List[str]],
    ):
        return _record_fetch(self, index, keys)
 def _record_fetch(self, index: int, keys: Union[str, List[str]]):
    if not getattr(self, "_data", None) and self._num_records == 0:
        raise RuntimeError("Store not loaded")
    if not 0 <= index < self._num_records:
        raise ValueError(
            f"Record index out of bounds: {index}, num_records={self._num_records}"
        )
    if isinstance(keys, str):
        return _fetch_record_key(self, keys, index)
    return {k: _fetch_record_key(self, k, index) for k in keys}
 def _fetch_record_key(self, key: str, index: int):
    offsets = self._offsets.get(key)
    if offsets:
        start = offsets[index]
        end = (
            offsets[index + 1]
            if index + 1 < len(offsets)
            else self._data[key][0].shape[0]
        )
        return self._data[key][0][start:end]
    return self._data[key][index]
 class StoreFactory(BaseFactory["Store"]):
    """Factory for creating Store instances by type name."""
@StoreFactory.register("h5")
 class H5Store(Store, Streamable, Recordable):
    """HDF5-based storage backend (pre-tokenized data).
    Each key is stored as a group of per-record datasets (``data_0``,
    ``data_1``, …).  Supports both access modes:
    - **Stream**: ``fetch(begin, end, key)`` and ``store[i]`` slice
      across concatenated records via ``_cum`` — used by SEQ/SFT.
    - **Record**: ``fetch_record(i, key)`` and ``store[i]`` (when
      ``window_size == 0``) index ``_data[key]`` directly — used by
      DPO/GRPO.
    """
    segments_are_records = True
    def __init__(
        self,
        window_size: int = 0,
        stride: Optional[int] = None,
    ):
        super().__init__(window_size=window_size, stride=stride)
    def load(self, path: str, **kwargs):
        self._normalize(load_h5(path))
@StoreFactory.register("bin")
 class MmapStore(Store, Streamable, Recordable):
    """Memory-mapped binary storage backend.
    Each key is a single .bin file backed by ``np.memmap(mode="r")``.
    No per-process memory duplication — all DataLoader workers share the
    same OS page-cache pages.
    Supports both access modes:
    - **Stream**: always available via :meth:`fetch`.
    - **Record** (``fetch_record(i, key)``): only when ``meta.json``
      contains per-record ``offsets`` (written via
      ``save_bin(..., record_keys=...)``).  Legacy bin files without
      offsets have ``num_records == 0`` and ``len(store)`` reflects the
      windowed sample count when ``window_size > 0``.
    ``segments_are_records`` is ``False`` here (bin segments are
    contiguous streams, not per-record) — record access is driven
    purely by ``_offsets``.
    """
    segments_are_records = False
    def __init__(
        self,
        window_size: int = 0,
        stride: Optional[int] = None,
    ):
        super().__init__(window_size=window_size, stride=stride)
        self._mmap_refs: List[Tensor] = []
    def load(self, path: str, **kwargs):
        self._mmap_refs = []
        root = Path(path)
        all_raw: Dict[str, List[Tensor]] = {}
        all_offsets: Dict[str, List[int]] = {}
        meta_paths = [
            Path(p) for p in glob.glob(str(root / "**" / "meta.json"), recursive=True)
        ]
        for meta_path in meta_paths:
            raw = load_bin(str(meta_path.parent))
            off = load_bin_offsets(str(meta_path.parent))
            for key, tensors in raw.items():
                if key not in all_raw:
                    all_raw[key] = []
                all_raw[key].extend(tensors)
            for key, o in off.items():
                if key not in all_offsets:
                    all_offsets[key] = []
                all_offsets[key].extend(o)
        if not meta_paths:
            raise FileNotFoundError(f"No meta.json found under {path}")
        self._normalize(all_raw, offsets=all_offsets or None)
        for tensors in self._data.values():
            self._mmap_refs.extend(tensors)
 class JsonlSource:
    """Read raw JSON records from a ``.jsonl`` file or directory.
    A thin reader used by :class:`JsonlStore` in processor mode — holds
    no tokenizer, performs no tokenisation, just yields dicts.
    """
    def __init__(self, path: str):
        self.path = Path(path)
        self._records: Optional[List[dict]] = None
    def load(self) -> List[dict]:
        if self._records is None:
            self._records = self._read(self.path)
        return self._records
    @staticmethod
    def _read(root: Path) -> List[dict]:
        if root.is_file():
            return JsonlSource._read_file(root)
        return JsonlSource._read_dir(root)
    @staticmethod
    def _read_file(path: Path) -> List[dict]:
        records: List[dict] = []
        with open(path, "r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                try:
                    records.append(json.loads(line))
                except json.JSONDecodeError:
                    logger.warning("Failed to parse JSON line in %s, skipping", path)
        return records
    @staticmethod
    def _read_dir(root: Path) -> List[dict]:
        records: List[dict] = []
        for jsonl_path in sorted(root.glob("*.jsonl")):
            records.extend(JsonlSource._read_file(jsonl_path))
        return records
@StoreFactory.register("jsonl")
 class JsonlStore(Store, Streamable, Recordable):
    """JSONL reader with two tokenisation modes.
    A JSONL dataset is a ``.jsonl`` file or a directory of ``*.jsonl``
    files plus (optionally) a ``dataset_config.json`` describing the
    tokenization pipeline.
    Two modes, selected at :meth:`load` time:
    - **Eager** (default): applies a :class:`TokenizeTransform` to every
      record at load time and registers per-key tensors via
      ``_normalize``.  Both ``fetch`` (stream) and ``fetch_record``
      (record) work.
    - **Lazy** (``processor=fn`` passed): keeps raw records and defers
      tokenisation to ``fetch_record``.  Only record access works —
      ``len(store)`` returns ``num_records``; stream primitives raise.
    """
    CONFIG_NAME = "dataset_config.json"
    segments_are_records = True
    def __init__(
        self,
        window_size: int = 0,
        stride: Optional[int] = None,
    ):
        super().__init__(window_size=window_size, stride=stride)
        self._source: Optional[JsonlSource] = None
        self._processor: Optional[Callable[[dict], Dict[str, Tensor]]] = None
        self._keys_cache: Optional[List[str]] = None
    def load(self, path: str, transform=None, processor=None, **kwargs):
        self._source = JsonlSource(path)
        records = self._source.load()
        if processor is not None:
            self._processor = processor
            self._num_records = len(records)
            return
        if transform is None:
            root = Path(path)
            config_path = root / self.CONFIG_NAME if root.is_dir() else None
            if config_path is None or not config_path.exists():
                raise FileNotFoundError(
                    f"JSONL dataset config not found. Expected "
                    f"{self.CONFIG_NAME} alongside *.jsonl files, pass an "
                    f"explicit transform, or pass processor= for lazy "
                    f"on-the-fly tokenisation."
                )
            transform = TokenizeTransform.from_config_file(str(config_path))
        transformed = transform.apply(records)
        self._normalize(transformed)
    @property
    def keys(self) -> List[str]:
        if self._processor is not None:
            if self._keys_cache is None and self._num_records > 0:
                sample = self._processor(self._source.load()[0])
                self._keys_cache = list(sample.keys())
            return self._keys_cache or []
        return list(self._data.keys())
    def fetch_record(self, index: int, keys: Union[str, List[str]]):
        if self._processor is not None:
            if not 0 <= index < self._num_records:
                raise ValueError(
                    f"Record index out of bounds: {index}, "
                    f"num_records={self._num_records}"
                )
            record = self._source.load()[index]
            data = self._processor(record)
            if isinstance(keys, str):
                return data[keys]
            return {k: data[k] for k in keys}
        return _record_fetch(self, index, keys)
    def fetch(self, begin: int, end: int, keys: Union[str, List[str]]):
        if self._processor is not None:
            raise RuntimeError(
                "JsonlStore in lazy (processor) mode does not support "
                "stream fetch(); use fetch_record() instead."
            )
        return _stream_fetch(self, begin, end, keys)
    def __getitem__(self, index: int) -> Dict[str, Tensor]:
        if self._processor is not None:
            return self.fetch_record(index, self._record_keys())
        return super().__getitem__(index)
@@ -0,0 +1,29 @@
 """CUDA attention kernel wrappers with torch fallback.
 Public API:
    - ``attn_decode`` — single-query decode attention
    - ``attn_prefill`` — multi-query prefill attention
    - ``attn_paged_decode`` — paged decode attention (direct page-table access)
 Interface (shared by all wrappers):
    causal_offset: -1 = non-causal; >=0 = absolute position of first Q token
    mask:       2D [batch, kv_len] or 3D [batch, q_len, kv_len] (bool, True = keep)
    scale:      0.0 = auto (1/sqrt(head_dim)); >0 = explicit
    layout:     "bhld" (default) or "blhd"
 Causal and mask can coexist — both are applied simultaneously.
 Each wrapper dispatches to its compiled CUDA kernel (``astrai.extension.attn_*``)
 when available, otherwise falls back to ``torch.nn.functional.scaled_dot_product_attention``.
 """
 from astrai.extension.loader import KERNEL_NAMES, is_available
 from astrai.extension.ops import attn_decode, attn_paged_decode, attn_prefill
 __all__ = [
    "attn_decode",
    "attn_paged_decode",
    "attn_prefill",
    "is_available",
    "KERNEL_NAMES",
 ]
@@ -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,246 @@
 """GQA attention wrapper functions — one entry point per compiled kernel.
 Each wrapper dispatches to its CUDA kernel (loaded in ``loader.py``) when
 available, otherwise falls back to ``torch`` SDPA.
 Interface (all functions):
    causal_offset: -1 = non-causal; >=0 = absolute position of first Q token
    mask:       2D [batch, kv_len] or 3D [batch, q_len, kv_len] (bool)
    scale:      0.0 = auto (1/sqrt(head_dim)); >0 = explicit
    layout:     "bhld" (default) or "blhd"
 Add new kernel wrappers here; split into per-variant files only if this file
 grows large.
 """
 import math
 import torch
 import torch.nn.functional as F
 from astrai.extension.loader import _available, _modules
 _LAYOUT_CODES: dict[str, int] = {"bhld": 0, "blhd": 1}
 def _parse_layout(layout: str | int) -> int:
    if isinstance(layout, int):
        return layout
    code = _LAYOUT_CODES.get(layout.lower())
    if code is None:
        raise ValueError(
            f"unknown layout '{layout}', expected one of {list(_LAYOUT_CODES)}"
        )
    return code
 def _to_bhld(t: torch.Tensor, layout: int) -> torch.Tensor:
    """Normalize to b h l d view. Zero-copy transpose if layout==1 (b l h d)."""
    if layout == 1:
        return t.transpose(1, 2)
    return t
 def _expand_kv_heads(
    k: torch.Tensor, v: torch.Tensor, q_head: int
 ) -> tuple[torch.Tensor, torch.Tensor]:
    """Expand K/V heads to match Q heads for GQA fallback."""
    kv_head = k.size(1)
    if kv_head == q_head:
        return k, v
    group = q_head // kv_head
    k = k.repeat_interleave(group, dim=1)
    v = v.repeat_interleave(group, dim=1)
    return k, v
 def _build_attn_mask(
    q: torch.Tensor,
    k: torch.Tensor,
    mask: torch.Tensor | None,
    causal_offset: int,
    scale: float,
 ) -> tuple[torch.Tensor | None, float]:
    """Build SDPA-compatible attn_mask + resolved scale.
    q and k must already be in b h l d layout.
    Causal and mask can coexist: causal sets -inf above the diagonal, mask
    sets -inf for padded positions. Both are OR'd into a single bool mask.
    """
    q_len = q.size(2)
    kv_len = k.size(2)
    head_dim = q.size(3)
    resolved_scale = scale if scale and scale > 0 else 1.0 / math.sqrt(head_dim)
    attn_mask = None
    if mask is not None:
        if mask.dim() == 2:
            # [batch, kv_len] → [batch, 1, 1, kv_len]
            attn_mask = mask[:, None, None, :]
        elif mask.dim() == 3:
            # [batch, q_len, kv_len] → [batch, 1, q_len, kv_len]
            attn_mask = mask[:, None, :, :]
        else:
            raise ValueError(f"mask must be 2D or 3D, got {mask.dim()}D")
    if causal_offset >= 0:
        batch = q.size(0)
        # q row i attends to kv cols 0..(causal_offset + i)
        q_idx = torch.arange(q_len, device=q.device).unsqueeze(1)  # [q_len, 1]
        kv_idx = torch.arange(kv_len, device=q.device).unsqueeze(0)  # [1, kv_len]
        causal_bool = kv_idx > (causal_offset + q_idx)  # True = masked out
        causal_mask = causal_bool.unsqueeze(0).expand(
            batch, -1, -1
        )  # [batch, q_len, kv_len]
        causal_mask = causal_mask[:, None, :, :]  # [batch, 1, q_len, kv_len]
        if attn_mask is not None:
            attn_mask = attn_mask | causal_mask
        else:
            attn_mask = causal_mask
    return attn_mask, resolved_scale
 def _torch_fallback(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    mask: torch.Tensor | None,
    causal_offset: int,
    scale: float,
    q_layout: int,
    kv_layout: int | None = None,
 ) -> torch.Tensor:
    """Reference attention via ``scaled_dot_product_attention``.
    q_layout / kv_layout: 0 = b h l d, 1 = b l h d.
    If kv_layout is None, uses q_layout (Q and K/V share the same layout).
    """
    if kv_layout is None:
        kv_layout = q_layout
    q = _to_bhld(q, q_layout)
    k = _to_bhld(k, kv_layout)
    v = _to_bhld(v, kv_layout)
    k, v = _expand_kv_heads(k, v, q.size(1))
    attn_mask, resolved_scale = _build_attn_mask(q, k, mask, causal_offset, scale)
    out = F.scaled_dot_product_attention(
        q, k, v, attn_mask=attn_mask, is_causal=False, scale=resolved_scale
    )
    # Restore Q's original layout
    if q_layout == 1:
        out = out.transpose(1, 2)
    return out
 def _gather_kv_from_pages(
    page_table: torch.Tensor,
    k_cache: torch.Tensor,
    v_cache: torch.Tensor,
    page_size: int,
    kv_len: int,
 ) -> tuple[torch.Tensor, torch.Tensor]:
    """Gather contiguous K/V from paged cache for torch SDPA fallback.
    Shapes:
        page_table : [batch, max_pages] (int64)
        k_cache    : [n_pages, page_size, n_kv_heads, head_dim]
        v_cache    : same as k_cache
    Returns:
        k, v : [batch, kv_len, n_kv_heads, head_dim]  (b l h d)
    """
    batch, max_pages = page_table.shape
    _, ps, n_kv_heads, head_dim = k_cache.shape
    if ps != page_size:
        raise ValueError(f"k_cache page_size mismatch: {ps} vs {page_size}")
    # Vectorized gather: build physical page + offset indices, then advanced-index
    positions = torch.arange(kv_len, device=page_table.device)
    logical_pages = positions // page_size  # [kv_len]
    page_offsets = positions % page_size  # [kv_len]
    phys_pages = page_table[:, logical_pages]  # [batch, kv_len]
    # k_cache[phys_pages, page_offsets] → [batch, kv_len, n_kv_heads, head_dim] (b l h d)
    k = k_cache[phys_pages, page_offsets]
    v = v_cache[phys_pages, page_offsets]
    return k, v
 def attn_decode(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    mask: torch.Tensor | None = None,
    causal_offset: int = -1,
    scale: float = 0.0,
    layout: str = "bhld",
 ) -> torch.Tensor:
    li = _parse_layout(layout)
    if _available["attn_decode"]:
        return _modules["attn_decode"].attn_decode(
            q,
            k,
            v,
            mask=mask,
            causal_offset=causal_offset,
            scale=scale,
            layout=li,
        )
    return _torch_fallback(q, k, v, mask, causal_offset, scale, q_layout=li)
 def attn_prefill(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    mask: torch.Tensor | None = None,
    causal_offset: int = -1,
    scale: float = 0.0,
    layout: str = "bhld",
 ) -> torch.Tensor:
    li = _parse_layout(layout)
    if _available["attn_prefill"]:
        return _modules["attn_prefill"].attn_prefill(
            q,
            k,
            v,
            mask=mask,
            causal_offset=causal_offset,
            scale=scale,
            layout=li,
        )
    return _torch_fallback(q, k, v, mask, causal_offset, scale, q_layout=li)
 def attn_paged_decode(
    q: torch.Tensor,
    page_table: torch.Tensor,
    k_cache: torch.Tensor,
    v_cache: torch.Tensor,
    page_size: int,
    kv_len: int,
    mask: torch.Tensor | None = None,
    causal_offset: int = -1,
    scale: float = 0.0,
    layout: str = "bhld",
 ) -> torch.Tensor:
    li = _parse_layout(layout)
    if _available["attn_paged_decode"]:
        return _modules["attn_paged_decode"].attn_paged_decode(
            q,
            page_table,
            k_cache,
            v_cache,
            page_size,
            kv_len,
            mask=mask,
            causal_offset=causal_offset,
            scale=scale,
            layout=li,
        )
    # Gathered K/V are always b l h d
    k, v = _gather_kv_from_pages(page_table, k_cache, v_cache, page_size, kv_len)
    return _torch_fallback(
        q, k, v, mask, causal_offset, scale, q_layout=li, kv_layout=1
    )
@@ -1,190 +1,144 @@
-"""Base factory class for extensible component registration."""
+"""Base factory with decorator-based registration and kwarg-filtered instantiation."""
 import inspect
 import sys
 from abc import ABC
-from typing import Callable, Dict, Generic, List, Optional, Tuple, Type, TypeVar
+from typing import (
    Callable,
    Dict,
    ForwardRef,
    Generic,
    List,
    Optional,
    Type,
    TypeVar,
    Union,
 )
 from typing import get_args as _get_args
 from typing import get_origin as _get_origin
 T = TypeVar("T")
-class Registry:
+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 list_registered(cls) -> list:
+    def get_component_class(cls, name: str) -> Type[T]:
-        """List all registered component names.
+        """Get the registered component class without instantiating it."""
        entry = cls._entries.get(name)
        if entry is None:
            raise ValueError(
                f"Unknown component: '{name}'. Supported types: {sorted(cls._entries)}"
            )
        return entry
-        Returns:
+    @classmethod
-            List of registered component names
+    def list_registered(cls) -> List[str]:
-        """
+        """List all registered component names."""
-        return cls._registry.list_names()
+        return sorted(cls._entries)
    @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,46 +1,105 @@
 """Inference module for continuous batching.
 Layers:
-  - engine.py:    Facade (InferenceEngine), Value Object (GenerationParams, GenerationRequest)
+  - core/:        Core inference loop (cache, executor, scheduler, task)
-  - scheduler.py: Continuous-batching loop, Task state machine, TaskStatus enum
+  - api/:         HTTP orchestration (ProtocolHandler, server)
-  - cache.py:     PagedCache (page-table-indirected KV cache with alloc/free)
+  - protocols/:   Response builders (OpenAI, Anthropic)
-  - sampling.py:  Strategy pattern (TemperatureStrategy, TopKStrategy, TopPStrategy)
+  - transport/:   SSE transport utilities
-  - server.py:    FastAPI HTTP server (OpenAI-compatible endpoints)
+  - engine.py:    Facade (InferenceEngine), Value Object (GenerationRequest)
  - sample.py:    Strategy pattern (TemperatureStrategy, TopKStrategy, TopPStrategy, FrequencyPenaltyStrategy)
 """
-from astrai.inference.engine import (
+from astrai.inference.api import (
-    GenerationParams,
+    AnthropicMessage,
-    GenerationRequest,
+    BaseToolParser,
-    InferenceEngine,
+    ChatCompletionRequest,
    ChatMessage,
    FunctionDef,
    GenContext,
    MessagesRequest,
    ProtocolHandler,
    SimpleJsonToolParser,
    StopChecker,
    ToolDef,
    ToolParserFactory,
    get_app,
    run_server,
 )
-from astrai.inference.sampling import (
+from astrai.inference.api.anthropic import AnthropicResponseBuilder
 from astrai.inference.api.openai import OpenAIResponseBuilder
 from astrai.inference.core import (
    STOP,
    Allocator,
    CacheView,
    ContiguousCache,
    ContiguousCacheView,
    Executor,
    InferenceScheduler,
    KVCache,
    PageCache,
    PageCacheView,
    PagePool,
    PrefixCache,
    Storage,
    Task,
    TaskManager,
    TaskStatus,
    TaskTable,
    page_hash,
 )
 from astrai.inference.engine import GenerationRequest, InferenceEngine
 from astrai.inference.sample import (
    BaseSamplingStrategy,
    FrequencyPenaltyStrategy,
    SamplingPipeline,
    TemperatureStrategy,
    TopKStrategy,
    TopPStrategy,
    sample,
 )
 from astrai.inference.scheduler import (
    InferenceScheduler,
    Task,
    TaskStatus,
 )
 __all__ = [
    # Engine / Requests
    "InferenceEngine",
    "GenerationRequest",
    "GenerationParams",
    # Scheduler
    "InferenceScheduler",
    "Executor",
    "STOP",
    "Task",
    "TaskManager",
    "TaskStatus",
-    # Sampling (Strategy pattern)
+    "Allocator",
    "CacheView",
    "KVCache",
    "ContiguousCache",
    "ContiguousCacheView",
    "PageCache",
    "PageCacheView",
    "PagePool",
    "PrefixCache",
    "Storage",
    "TaskTable",
    "page_hash",
    "sample",
    "BaseSamplingStrategy",
    "TemperatureStrategy",
    "TopKStrategy",
    "TopPStrategy",
    "FrequencyPenaltyStrategy",
    "SamplingPipeline",
    "ProtocolHandler",
    "StopChecker",
    "GenContext",
    "BaseToolParser",
    "SimpleJsonToolParser",
    "ToolParserFactory",
    "OpenAIResponseBuilder",
    "AnthropicResponseBuilder",
    "ChatMessage",
    "ChatCompletionRequest",
    "FunctionDef",
    "ToolDef",
    "AnthropicMessage",
    "MessagesRequest",
    "get_app",
    "run_server",
 ]
@@ -0,0 +1,39 @@
 """Inference API: protocol handler, stop checker, tool parsers, and FastAPI server.
 ``app`` is no longer a module-level global. Use :func:`get_app` to access the
 lazy singleton FastAPI instance.
 """
 from astrai.inference.api.protocol import GenContext, ProtocolHandler, StopChecker
 from astrai.inference.api.server import (
    AnthropicMessage,
    ChatCompletionRequest,
    ChatMessage,
    FunctionDef,
    MessagesRequest,
    ToolDef,
    get_app,
    run_server,
 )
 from astrai.inference.api.tool_parser import (
    BaseToolParser,
    SimpleJsonToolParser,
    ToolParserFactory,
 )
 __all__ = [
    "ProtocolHandler",
    "StopChecker",
    "GenContext",
    "BaseToolParser",
    "SimpleJsonToolParser",
    "ToolParserFactory",
    "AnthropicMessage",
    "ChatCompletionRequest",
    "ChatMessage",
    "FunctionDef",
    "ToolDef",
    "MessagesRequest",
    "get_app",
    "run_server",
 ]
@@ -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,277 @@
 """OpenAI chat completion response builder."""
 import logging
 import time
 import uuid
 from typing import Any, Dict, List, Optional, Tuple, Union
 from pydantic import BaseModel
 from astrai.inference.api.protocol import (
    GenContext,
    ResponseBuilder,
    StopInfo,
    sse_event,
 )
 from astrai.inference.api.tool_parser import BaseToolParser, ToolParserFactory
 from astrai.inference.engine import InferenceEngine
 logger = logging.getLogger(__name__)
 _UNSUPPORTED_PARAMS = (
    "n",
    "presence_penalty",
    "logit_bias",
    "user",
 )
 def _resolve_tool_choice(
    request: BaseModel,
 ) -> Union[str, Dict[str, Any]]:
    tc = getattr(request, "tool_choice", None)
    if tc is None:
        return "auto"
    if isinstance(tc, str):
        return tc
    if isinstance(tc, dict):
        return tc
    return "auto"
 def _resolve_tools(request: BaseModel) -> Optional[List[Dict[str, Any]]]:
    raw = getattr(request, "tools", None)
    if not raw:
        return None
    if isinstance(raw, list):
        return [t.model_dump() if hasattr(t, "model_dump") else t for t in raw]
    return None
 class OpenAIResponseBuilder(ResponseBuilder):
    def prepare(
        self, request: BaseModel, engine: InferenceEngine
    ) -> Tuple[str, GenContext, List[str]]:
        messages = [{"role": m.role, "content": m.content} for m in request.messages]
        tools = _resolve_tools(request)
        prompt = engine.tokenizer.apply_chat_template(
            messages, tokenize=False, tools=tools or []
        )
        self._resp_id = f"chatcmpl-{uuid.uuid4().hex[:12]}"
        self._model = request.model
        for param in _UNSUPPORTED_PARAMS:
            value = getattr(request, param, None)
            fields = getattr(type(request), "model_fields", {})
            default = fields[param].default if param in fields else None
            if value is not None and value != default:
                logger.warning(
                    "ChatCompletionRequest param '%s'=%r is not supported"
                    " and will be ignored",
                    param,
                    value,
                )
        self._parser: Optional[BaseToolParser] = None
        if tools:
            tool_choice = _resolve_tool_choice(request)
            self._parser = ToolParserFactory.create(
                "simple_json", tools=tools, tool_choice=tool_choice
            )
        self._content_started = False
        ctx = GenContext(
            resp_id=self._resp_id,
            created=int(time.time()),
            model=self._model,
        )
        stop = request.stop
        stop_sequences = (
            [] if stop is None else [stop] if isinstance(stop, str) else stop
        )
        return prompt, ctx, stop_sequences
    def format_stream_start(self, ctx: GenContext) -> List[str]:
        return [
            sse_event(
                {
                    "id": self._resp_id,
                    "object": "chat.completion.chunk",
                    "created": ctx.created,
                    "model": self._model,
                    "choices": [
                        {
                            "index": 0,
                            "delta": {"role": "assistant"},
                            "finish_reason": None,
                        }
                    ],
                }
            )
        ]
    def format_chunk(self, token: str, **kwargs) -> List[str]:
        body = kwargs.get("body", "")
        if self._parser is not None:
            return self._format_tool_chunk(body, **kwargs)
        return [
            sse_event(
                {
                    "id": self._resp_id,
                    "object": "chat.completion.chunk",
                    "created": 0,
                    "model": self._model,
                    "choices": [
                        {
                            "index": 0,
                            "delta": {"content": token},
                            "finish_reason": None,
                        }
                    ],
                }
            )
        ]
    def _format_tool_chunk(self, body: str, **kwargs) -> List[str]:
        deltas = self._parser.feed(
            body,
            current_token_ids=kwargs.get("current_token_ids"),
            delta_token_ids=kwargs.get("delta_token_ids"),
        )
        events: List[str] = []
        for d in deltas:
            if "content" in d:
                if not self._content_started:
                    events.append(self._role_chunk())
                    self._content_started = True
                events.append(
                    sse_event(
                        {
                            "id": self._resp_id,
                            "object": "chat.completion.chunk",
                            "created": 0,
                            "model": self._model,
                            "choices": [
                                {
                                    "index": 0,
                                    "delta": {"content": d["content"]},
                                    "finish_reason": None,
                                }
                            ],
                        }
                    )
                )
            elif "tool_calls" in d:
                if not self._content_started:
                    events.append(self._role_chunk())
                    self._content_started = True
                events.append(
                    sse_event(
                        {
                            "id": self._resp_id,
                            "object": "chat.completion.chunk",
                            "created": 0,
                            "model": self._model,
                            "choices": [
                                {
                                    "index": 0,
                                    "delta": {"tool_calls": d["tool_calls"]},
                                    "finish_reason": None,
                                }
                            ],
                        }
                    )
                )
        return events
    def _role_chunk(self) -> str:
        return sse_event(
            {
                "id": self._resp_id,
                "object": "chat.completion.chunk",
                "created": 0,
                "model": self._model,
                "choices": [
                    {
                        "index": 0,
                        "delta": {"role": "assistant"},
                        "finish_reason": None,
                    }
                ],
            }
        )
    def format_stream_end(self, ctx: GenContext, stop: StopInfo) -> List[str]:
        finish_reason = "stop"
        if self._parser is not None and self._parser.has_tool_calls:
            finish_reason = "tool_calls"
        return [
            sse_event(
                {
                    "id": self._resp_id,
                    "object": "chat.completion.chunk",
                    "created": ctx.created,
                    "model": self._model,
                    "choices": [
                        {"index": 0, "delta": {}, "finish_reason": finish_reason}
                    ],
                }
            ),
            sse_event(
                {
                    "prompt_tokens": ctx.prompt_tokens,
                    "completion_tokens": ctx.completion_tokens,
                    "total_tokens": ctx.prompt_tokens + ctx.completion_tokens,
                }
            ),
        ]
    def format_response(
        self, ctx: GenContext, content: str, stop: StopInfo
    ) -> Dict[str, Any]:
        if self._parser is not None:
            parsed = self._parser.parse_complete(content)
            if parsed and parsed.get("tool_calls"):
                return {
                    "id": self._resp_id,
                    "object": "chat.completion",
                    "created": ctx.created,
                    "model": self._model,
                    "choices": [
                        {
                            "index": 0,
                            "message": {
                                "role": "assistant",
                                "content": parsed.get("content"),
                                "tool_calls": parsed["tool_calls"],
                            },
                            "finish_reason": "tool_calls",
                        }
                    ],
                    "usage": {
                        "prompt_tokens": ctx.prompt_tokens,
                        "completion_tokens": ctx.completion_tokens,
                        "total_tokens": ctx.prompt_tokens + ctx.completion_tokens,
                    },
                }
        return {
            "id": self._resp_id,
            "object": "chat.completion",
            "created": ctx.created,
            "model": self._model,
            "choices": [
                {
                    "index": 0,
                    "message": {"role": "assistant", "content": content},
                    "finish_reason": "stop",
                }
            ],
            "usage": {
                "prompt_tokens": ctx.prompt_tokens,
                "completion_tokens": ctx.completion_tokens,
                "total_tokens": ctx.prompt_tokens + ctx.completion_tokens,
            },
        }
@@ -0,0 +1,200 @@
 """Orchestration layer: ProtocolHandler, StopChecker, GenContext, StopInfo, ResponseBuilder, SSE utils.
 ProtocolHandler orchestrates the async generation loop and delegates
 protocol-specific formatting to a ResponseBuilder.
 """
 import json
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from typing import Any, AsyncGenerator, Dict, List, Optional, Tuple, Union
 from fastapi.responses import StreamingResponse
 from pydantic import BaseModel
 from astrai.inference.engine import InferenceEngine
 def sse_event(data: Dict[str, Any], event: Optional[str] = None) -> str:
    lines: List[str] = []
    if event:
        lines.append(f"event: {event}")
    lines.append(f"data: {json.dumps(data, ensure_ascii=False)}")
    lines.append("")
    return "\n".join(lines)
 def sse_done() -> str:
    return "data: [DONE]\n\n"
@dataclass
 class GenContext:
    """Per-generation metadata passed to builder format methods."""
    resp_id: str
    created: int
    model: str
    prompt_tokens: int = 0
    completion_tokens: int = 0
@dataclass
 class StopInfo:
    """Stop-check result passed to format_stream_end / format_response."""
    matched: Optional[str] = None
    body: str = ""
    yielded: str = ""
 class StopChecker:
    """Scans accumulated text for stop sequence matches."""
    def __init__(self, sequences: List[str]):
        self._sequences = [s for s in sequences if s]
    def check(self, text: str) -> Optional[str]:
        for seq in self._sequences:
            if seq in text:
                return seq
        return None
 class ResponseBuilder(ABC):
    """Interface for protocol-specific response formatting.
    A new protocol requires one concrete builder implementing 5 methods.
    """
    @abstractmethod
    def prepare(
        self, request: BaseModel, engine: InferenceEngine
    ) -> Tuple[str, GenContext, List[str]]:
        """Return (prompt, ctx, stop_sequences) for a generation request."""
    @abstractmethod
    def format_stream_start(self, ctx: GenContext) -> List[str]:
        """SSE events that open the stream."""
    @abstractmethod
    def format_chunk(self, token: str, **kwargs) -> List[str]:
        """SSE events for a single generated token.
        ``body`` (the full accumulated text so far) is always provided
        as a keyword argument. Additional keyword arguments such as
        ``current_token_ids`` and ``delta_token_ids`` may be included
        for tool parsers that need token-level information.
        Returns a list of SSE event strings (may be empty).
        """
    @abstractmethod
    def format_stream_end(self, ctx: GenContext, stop: StopInfo) -> List[str]:
        """SSE events that close the stream."""
    @abstractmethod
    def format_response(
        self, ctx: GenContext, content: str, stop: StopInfo
    ) -> Dict[str, Any]:
        """JSON response body for non-streaming mode."""
 class ProtocolHandler:
    """Orchestrates the generation loop, delegates formatting to a builder.
    Usage::
        handler = ProtocolHandler(request, engine, OpenAIResponseBuilder())
        response = await handler.handle()
    """
    def __init__(
        self, request: BaseModel, engine: InferenceEngine, builder: ResponseBuilder
    ):
        self.request = request
        self.engine = engine
        self.builder = builder
    async def handle(self) -> Union[StreamingResponse, Dict[str, Any]]:
        prompt, ctx, stop_sequences = self.builder.prepare(self.request, self.engine)
        ctx.prompt_tokens = len(self.engine.tokenizer.encode(prompt))
        agen = self.engine.generate_async(
            prompt=prompt,
            max_tokens=self.request.max_tokens,
            temperature=self.request.temperature,
            top_p=self.request.top_p,
            top_k=self.request.top_k,
            frequency_penalty=getattr(self.request, "frequency_penalty", 0.0),
        )
        if self.request.stream:
            return self._handle_stream(agen, ctx, stop_sequences)
        else:
            return await self._handle_non_stream(agen, ctx, stop_sequences)
    def _handle_stream(
        self, agen: AsyncGenerator, ctx: GenContext, stop_sequences: List[str]
    ) -> StreamingResponse:
        checker = StopChecker(stop_sequences)
        async def event_stream():
            for event in self.builder.format_stream_start(ctx):
                yield event
            body = ""
            yielded = ""
            matched = None
            token_ids: List[int] = []
            async for token in agen:
                body += token
                new_ids = self.engine.tokenizer.encode(token)
                token_ids.extend(new_ids)
                matched = checker.check(body)
                if matched:
                    break
                ctx.completion_tokens += 1
                for event in self.builder.format_chunk(
                    token,
                    body=body,
                    current_token_ids=token_ids,
                    delta_token_ids=new_ids,
                ):
                    yield event
                yielded += token
            stop = StopInfo(matched=matched, body=body, yielded=yielded)
            for event in self.builder.format_stream_end(ctx, stop):
                yield event
            yield sse_done()
        return StreamingResponse(
            event_stream(),
            media_type="text/event-stream",
            headers={"Cache-Control": "no-cache", "Connection": "keep-alive"},
        )
    async def _handle_non_stream(
        self, agen: AsyncGenerator, ctx: GenContext, stop_sequences: List[str]
    ) -> Dict[str, Any]:
        checker = StopChecker(stop_sequences)
        chunks: List[str] = []
        body = ""
        matched = None
        async for token in agen:
            chunks.append(token)
            body += token
            matched = checker.check(body)
            if matched:
                break
            ctx.completion_tokens += 1
        content = "".join(chunks)
        stop = StopInfo(matched=matched, body=body)
        return self.builder.format_response(ctx, content, stop)
@@ -0,0 +1,202 @@
 """
 OpenAI / Anthropic-compatible chat completion server backed by continuous-batching inference.
 Protocol-specific formatting is delegated to ``astrai.inference.protocol``.
 This module owns the FastAPI app, request/response schemas, and dependency wiring.
 ``app`` is lazily constructed — importing this module does NOT create a FastAPI instance.
 Use :func:`get_app` to access the singleton.
 """
 import logging
 from contextlib import asynccontextmanager
 from pathlib import Path
 from typing import Any, Dict, List, Optional, Union
 import torch
 import uvicorn
 from fastapi import APIRouter, FastAPI, HTTPException
 from pydantic import BaseModel, Field
 from astrai.inference.api.anthropic import AnthropicResponseBuilder
 from astrai.inference.api.openai import OpenAIResponseBuilder
 from astrai.inference.api.protocol import ProtocolHandler
 from astrai.inference.engine import InferenceEngine
 from astrai.model import AutoModel
 from astrai.tokenize import AutoTokenizer
 logger = logging.getLogger(__name__)
 _app_instance: Optional[FastAPI] = None
 class ChatMessage(BaseModel):
    role: str
    content: Optional[str] = None
    tool_calls: Optional[List[Dict[str, Any]]] = None
    tool_call_id: Optional[str] = None
 class FunctionDef(BaseModel):
    name: str
    description: Optional[str] = None
    parameters: Optional[Dict[str, Any]] = None
 class ToolDef(BaseModel):
    type: str = "function"
    function: FunctionDef
 class ChatCompletionRequest(BaseModel):
    """OpenAI Chat Completion API request body."""
    model: str = "astrai"
    messages: List[ChatMessage]
    temperature: Optional[float] = Field(default=1.0, ge=0.0, le=2.0)
    top_p: Optional[float] = Field(default=1.0, ge=0.0, le=1.0)
    top_k: Optional[int] = Field(default=50, ge=1)
    stream: Optional[bool] = False
    stop: Optional[Union[str, List[str]]] = None
    max_tokens: Optional[int] = Field(default=2048, ge=1)
    n: Optional[int] = Field(default=1, ge=1)
    presence_penalty: Optional[float] = Field(default=0.0, ge=-2.0, le=2.0)
    frequency_penalty: Optional[float] = Field(default=0.0, ge=-2.0, le=2.0)
    logit_bias: Optional[Dict[int, float]] = None
    user: Optional[str] = None
    tools: Optional[List[ToolDef]] = None
    tool_choice: Optional[Union[str, Dict[str, Any]]] = "auto"
 class AnthropicMessage(BaseModel):
    role: str
    content: Union[str, List[Dict[str, Any]]]
 class MessagesRequest(BaseModel):
    """Anthropic Messages API request body."""
    model: str = "astrai"
    max_tokens: int = Field(default=1024, ge=1)
    messages: List[AnthropicMessage]
    system: Optional[str] = None
    temperature: Optional[float] = Field(default=1.0, ge=0.0, le=2.0)
    top_p: Optional[float] = Field(default=1.0, ge=0.0, le=1.0)
    top_k: Optional[int] = Field(default=50, ge=1)
    stream: Optional[bool] = False
    stop_sequences: Optional[List[str]] = None
@asynccontextmanager
 async def lifespan(app: FastAPI):
    config = app.state.server_config
    if not config.get("_test", False):
        try:
            app.state.engine = _create_engine(**config)
        except Exception as e:
            logger.error(f"Failed to load model: {e}")
            raise
    yield
    if app.state.engine:
        app.state.engine.shutdown()
        logger.info("Inference engine shutdown complete")
 router = APIRouter()
 def _create_engine(
    param_path: Path,
    device: str = "cuda",
    dtype: torch.dtype = torch.bfloat16,
    max_batch_size: int = 16,
 ) -> InferenceEngine:
    if not param_path.exists():
        raise FileNotFoundError(f"Parameter directory not found: {param_path}")
    tokenizer = AutoTokenizer.from_pretrained(param_path)
    model = AutoModel.from_pretrained(param_path)
    model.to(device=device, dtype=dtype)
    logger.info(f"Model loaded on {device} with dtype {dtype}")
    engine = InferenceEngine(
        model=model,
        tokenizer=tokenizer,
        max_batch_size=max_batch_size,
    )
    logger.info(f"Inference engine initialized with max_batch_size={max_batch_size}")
    return engine
 def get_app() -> FastAPI:
    """Return the singleton FastAPI instance (lazily created on first call)."""
    global _app_instance
    if _app_instance is None:
        _app_instance = FastAPI(
            title="AstrAI Inference Server",
            version="0.2.0",
            lifespan=lifespan,
        )
        _app_instance.include_router(router)
        _app_instance.state.server_config = {}
        _app_instance.state.engine = None
    return _app_instance
 def _get_engine() -> InferenceEngine:
    engine = get_app().state.engine
    if engine is None:
        raise HTTPException(status_code=503, detail="Engine not initialized")
    return engine
@router.get("/health")
 async def health():
    app = get_app()
    return {
        "status": "ok",
        "model_loaded": app.state.engine is not None,
    }
@router.get("/stats")
 async def get_stats():
    return _get_engine().get_stats()
@router.post("/v1/chat/completions")
 async def chat_completion(request: ChatCompletionRequest):
    engine = _get_engine()
    handler = ProtocolHandler(request, engine, OpenAIResponseBuilder())
    return await handler.handle()
@router.post("/v1/messages")
 async def create_message(request: MessagesRequest):
    engine = _get_engine()
    handler = ProtocolHandler(request, engine, AnthropicResponseBuilder())
    return await handler.handle()
 def run_server(
    param_path: Path,
    host: str = "0.0.0.0",
    port: int = 8000,
    reload: bool = False,
    device: str = "cuda",
    dtype: torch.dtype = torch.bfloat16,
    max_batch_size: int = 16,
 ):
    app = get_app()
    app.state.server_config = {
        "device": device,
        "dtype": dtype,
        "param_path": param_path,
        "max_batch_size": max_batch_size,
    }
    uvicorn.run(
        app,
        host=host,
        port=port,
        reload=reload,
    )
@@ -0,0 +1,344 @@
 """Tool call parsers for extracting structured tool calls from model output.
 Patterned after vLLM's ToolParser abstraction. Each parser knows how to
 detect and incrementally extract tool calls from raw generated text.
 Subclasses may optionally consume ``token_ids`` for token-level parsing
 (e.g. Harmony / VLM-style parsers).
 """
 import json
 import re
 import uuid
 from abc import ABC, abstractmethod
 from typing import Dict, List, Optional
 from astrai.factory import BaseFactory
 class BaseToolParser(ABC):
    """Abstract tool call parser — one instance per request.
    Maintains streaming state internally so that each call to :meth:`feed`
    can diff against previously emitted content.
    Parameters
    ----------
    tools : list of dict, optional
        Tool definitions from the request.
    tool_choice : str
        ``"auto"`` / ``"required"`` / ``"none"`` or a named tool choice
        dict.
    """
    def __init__(self, tools: Optional[List[Dict]] = None, tool_choice: str = "auto"):
        self.tools = tools or []
        self.tool_choice = tool_choice
    @abstractmethod
    def feed(
        self,
        body: str,
        current_token_ids: Optional[List[int]] = None,
        delta_token_ids: Optional[List[int]] = None,
    ) -> List[Dict]:
        """Feed the *full* accumulated text each step.
        Returns a list of delta dicts to emit. Each delta is one of:
        - ``{"content": "text"}``      — plain text delta
        - ``{"tool_calls": [...]}``    — tool-call delta (OpenAI format)
        Returns an empty list when nothing new should be emitted.
        Parameters
        ----------
        body : str
            The complete accumulated generated text so far.
        current_token_ids : list of int, optional
            All token IDs decoded into *body* (cumulative).
        delta_token_ids : list of int, optional
            Only the token IDs for this chunk.
        """
    @abstractmethod
    def parse_complete(self, body: str) -> Optional[Dict]:
        """Parse the *complete* generated text after generation ends.
        Returns ``None`` when no tool calls were found, otherwise a dict
        with ``content`` (str or None) and ``tool_calls`` (list of dicts).
        """
    @property
    @abstractmethod
    def has_tool_calls(self) -> bool:
        """True if the parser detected at least one tool call in the stream."""
 class ToolParserFactory(BaseFactory["BaseToolParser"]):
    pass
 _TOOL_CALL_HEAD_RE = re.compile(r'\{\s*"name"\s*:')
 def _scan_json(text: str, start: int = 0):
    """Scan for a complete JSON object starting at *start*.
    Returns ``(end, complete)`` where *end* is one-past the closing
    brace (or ``len(text)`` if unclosed), and *complete* is a bool.
    """
    depth = 0
    in_string = False
    escape = False
    for i in range(start, len(text)):
        c = text[i]
        if escape:
            escape = False
            continue
        if c == "\\":
            escape = True
            continue
        if c == '"':
            in_string = not in_string
            continue
        if in_string:
            continue
        if c == "{":
            depth += 1
        elif c == "}":
            depth -= 1
            if depth == 0:
                return i + 1, True
    return len(text), False
 def _parse_tool_call_json(json_str: str, complete: bool):
    """Extract *name* and *arguments* from a tool-call JSON string.
    Returns ``(name, args, valid)``.
    """
    if complete:
        try:
            obj = json.loads(json_str)
        except json.JSONDecodeError:
            return None, "", False
        name = obj.get("name")
        if not isinstance(name, str) or not name:
            return None, "", False
        args = obj.get("arguments")
        if isinstance(args, dict):
            if not args:
                args = ""
            else:
                args = json.dumps(args, ensure_ascii=False)
                args = args[1:-1].rstrip()
        elif isinstance(args, list):
            args = json.dumps(args, ensure_ascii=False) if args else ""
        elif isinstance(args, str):
            pass
        else:
            args = str(args) if args is not None else ""
        return name, args, True
    name_match = re.search(r'"name"\s*:\s*"([^"]*)"', json_str)
    if not name_match:
        return None, "", False
    name = name_match.group(1)
    args_match = re.search(r'"arguments"\s*:\s*(.*)', json_str, re.DOTALL)
    if not args_match:
        return name, "", True
    raw = args_match.group(1).rstrip()
    if raw.startswith("{"):
        inner = raw[1:].rstrip()
        if inner.endswith("}"):
            inner = inner[:-1].rstrip()
        raw = inner
    return name, raw, True
 def _find_tool_calls(text: str, start_pos: int = 0):
    """Find all complete ``{...}`` tool-call objects in *text*.
    Returns a list of dicts with keys *start*, *end*, *name*, *args*,
    *complete*.
    """
    results = []
    pos = start_pos
    while True:
        brace = text.find("{", pos)
        if brace == -1:
            break
        end, complete = _scan_json(text, brace)
        if not complete:
            break
        json_str = text[brace:end]
        name, args, valid = _parse_tool_call_json(json_str, complete=True)
        if not valid or name is None:
            pos = end
            continue
        results.append(
            {
                "start": brace,
                "end": end,
                "name": name,
                "args": args,
                "complete": True,
            }
        )
        pos = end
    return results
 def _find_partial_tool_call(text: str, start_pos: int = 0):
    """Find one incomplete (still-generating) tool-call JSON object."""
    brace = text.find("{", start_pos)
    if brace == -1:
        return None
    json_str = text[brace:]
    if '"name"' not in json_str:
        return None
    name, args, valid = _parse_tool_call_json(json_str, complete=False)
    if not valid or name is None:
        return None
    return {
        "start": brace,
        "name": name,
        "args": args,
        "complete": False,
    }
@ToolParserFactory.register("simple_json")
 class SimpleJsonToolParser(BaseToolParser):
    """Parser for models that output tool calls as plain JSON objects.
    Detects ``{"name": "<func>", "arguments": {...}}`` anywhere in the
    generated text.  Handles single and (non-overlapping) multiple tool
    calls.  Text preceding the first tool call is emitted as plain
    ``content`` deltas.
    """
    def __init__(self, tools=None, tool_choice="auto"):
        super().__init__(tools, tool_choice)
        self._emitted_content_len = 0
        self._tc_state: List[Dict] = []
        self._has_tool_calls = False
    # -------------------------------------------------------------- feed
    def feed(
        self,
        body: str,
        current_token_ids: Optional[List[int]] = None,
        delta_token_ids: Optional[List[int]] = None,
    ) -> List[Dict]:
        deltas: List[Dict] = []
        completed = _find_tool_calls(body)
        if not completed:
            partial = _find_partial_tool_call(body)
            if not partial:
                return self._emit_plain_content(body, deltas)
            all_tcs = [partial]
        else:
            all_tcs = completed
            partial = _find_partial_tool_call(body, completed[-1]["end"])
            if partial:
                all_tcs = completed + [partial]
        first_start = all_tcs[0]["start"]
        if first_start > self._emitted_content_len:
            content = body[self._emitted_content_len : first_start]
            self._emitted_content_len = first_start
            if content:
                deltas.append({"content": content})
        for i, tc in enumerate(all_tcs):
            if i >= len(self._tc_state):
                self._tc_state.append(
                    {
                        "id": f"call_{uuid.uuid4().hex[:12]}",
                        "name_emitted": False,
                        "args_emitted_len": 0,
                    }
                )
                self._has_tool_calls = True
            st = self._tc_state[i]
            if not st["name_emitted"]:
                st["name_emitted"] = True
                deltas.append(
                    {
                        "tool_calls": [
                            {
                                "index": i,
                                "id": st["id"],
                                "type": "function",
                                "function": {"name": tc["name"], "arguments": ""},
                            }
                        ]
                    }
                )
            new_args = tc["args"]
            if len(new_args) > st["args_emitted_len"]:
                diff = new_args[st["args_emitted_len"] :]
                st["args_emitted_len"] = len(new_args)
                deltas.append(
                    {
                        "tool_calls": [
                            {
                                "index": i,
                                "function": {"arguments": diff},
                            }
                        ]
                    }
                )
        return deltas
    def _emit_plain_content(self, body: str, deltas: List[Dict]) -> List[Dict]:
        new_content = body[self._emitted_content_len :]
        if new_content:
            self._emitted_content_len = len(body)
            deltas.append({"content": new_content})
        return deltas
    # -------------------------------------------------------- complete
    def parse_complete(self, body: str) -> Optional[Dict]:
        completed = _find_tool_calls(body)
        if not completed:
            return None
        content = body[: completed[0]["start"]].strip() or None
        tool_calls = []
        for i, tc in enumerate(completed):
            tool_calls.append(
                {
                    "id": f"call_{uuid.uuid4().hex[:12]}",
                    "type": "function",
                    "function": {
                        "name": tc["name"],
                        "arguments": tc["args"],
                    },
                }
            )
        return {"content": content, "tool_calls": tool_calls}
    @property
    def has_tool_calls(self) -> bool:
        return self._has_tool_calls
@@ -1,174 +0,0 @@
 """Page-based KV cache with page-table-indirected read/write.
 Provides:
  - PagedCache: paged KV cache combining page pool and tensor storage.
 """
 from typing import Dict, List, Tuple
 import torch
 from torch import Tensor
 STOP = object()
 def page_hash(token_ids: List[int], page_idx: int, page_size: int) -> int:
    start = page_idx * page_size
    end = min(start + page_size, len(token_ids))
    h = 0
    for i in range(start, end):
        h = (h * 31 + token_ids[i]) & 0xFFFFFFFFFFFFFFFF
    return h
 class PagedCache:
    """Paged KV cache with page-table-indirected read/write.
    Combines:
      - Page pool (ref-counted alloc/free via bitmask)
      - KV tensor storage (k_cache, v_cache)
      - Prefix-cache hash lookup (page_content_hash -> physical_page_idx)
    Call :meth:`bind` to obtain a batch view for the attention layers.
    """
    def __init__(
        self,
        n_layers: int,
        n_pages: int,
        page_size: int,
        n_kv_heads: int,
        head_dim: int,
        device: torch.device,
        dtype: torch.dtype,
    ):
        self.page_size = page_size
        self._free_mask = (1 << n_pages) - 1
        self._refs: List[int] = [0] * n_pages
        self.k_cache = torch.empty(
            (n_layers, n_pages, page_size, n_kv_heads, head_dim),
            device=device,
            dtype=dtype,
        )
        self.v_cache = torch.empty(
            (n_layers, n_pages, page_size, n_kv_heads, head_dim),
            device=device,
            dtype=dtype,
        )
        self._page_to_hash: Dict[int, int] = {}
        self._hash_to_page: Dict[int, int] = {}
    def record_page(
        self, page_idx: int, token_ids: List[int], logical_page_idx: int
    ) -> None:
        h = page_hash(token_ids, logical_page_idx, self.page_size)
        old_h = self._page_to_hash.pop(page_idx, None)
        if old_h is not None:
            self._hash_to_page.pop(old_h, None)
        self._page_to_hash[page_idx] = h
        self._hash_to_page[h] = page_idx
    def lookup_prefix(self, token_ids: List[int]) -> List[int]:
        full_pages = len(token_ids) // self.page_size
        hits: List[int] = []
        for i in range(full_pages):
            h = page_hash(token_ids, i, self.page_size)
            p = self._hash_to_page.get(h)
            if p is None:
                break
            hits.append(p)
        return hits
    def inc_ref(self, idx: int) -> None:
        self._refs[idx] += 1
    def alloc(self) -> int:
        lsb = self._free_mask & -self._free_mask
        if lsb == 0:
            return -1
        idx = lsb.bit_length() - 1
        self._free_mask ^= lsb
        self._refs[idx] = 1
        return idx
    def alloc_n(self, n: int) -> List[int]:
        pages = [self.alloc() for _ in range(n)]
        if any(p < 0 for p in pages):
            for p in pages:
                if p >= 0:
                    self.free(p)
            return []
        return pages
    def free(self, idx: int) -> None:
        self._refs[idx] -= 1
        if self._refs[idx] == 0:
            self._free_mask |= 1 << idx
            h = self._page_to_hash.pop(idx, None)
            if h is not None:
                self._hash_to_page.pop(h, None)
    def bind(self, page_table: Tensor, total_len: int = 0) -> "CacheView":
        return CacheView(self, page_table, total_len)
    def write(
        self, layer_id: int, page_table: Tensor, start_pos: int, k: Tensor, v: Tensor
    ) -> None:
        seq_len = k.size(1)
        if seq_len == 0:
            return
        page_size = self.page_size
        written = 0
        first_page = start_pos // page_size
        last_page = (start_pos + seq_len - 1) // page_size
        for pi in range(first_page, last_page + 1):
            phys_pages = page_table[:, pi]
            page_start = pi * page_size
            write_start = max(page_start, start_pos)
            write_end = min(page_start + page_size, start_pos + seq_len)
            offset = write_start - page_start
            chunk = write_end - write_start
            self.k_cache[layer_id, phys_pages, offset : offset + chunk] = k[
                :, written : written + chunk
            ]
            self.v_cache[layer_id, phys_pages, offset : offset + chunk] = v[
                :, written : written + chunk
            ]
            written += chunk
    def gather(self, layer_id: int, page_table: Tensor) -> Tuple[Tensor, Tensor]:
        k_parts, v_parts = [], []
        for pi in range(page_table.size(1)):
            phys_pages = page_table[:, pi]
            if not (phys_pages >= 0).any():
                break
            k_parts.append(self.k_cache[layer_id, phys_pages])
            v_parts.append(self.v_cache[layer_id, phys_pages])
        k = torch.cat(k_parts, dim=1)
        v = torch.cat(v_parts, dim=1)
        return k, v
 class CacheView:
    """Per-batch view that bundles PagedCache + page_table + total_len.
    Attention layers receive this as ``paged_cache`` and only see
    ``write()`` / ``gather()``, never raw page tables or length params.
    """
    __slots__ = ("_cache", "_page_table", "_total_len")
    def __init__(self, cache: PagedCache, page_table: Tensor, total_len: int = 0):
        self._cache = cache
        self._page_table = page_table
        self._total_len = total_len
    def write(self, layer_id: int, start_pos: int, k: Tensor, v: Tensor) -> None:
        self._cache.write(layer_id, self._page_table, start_pos, k, v)
    def gather(self, layer_id: int) -> Tuple[Tensor, Tensor]:
        k, v = self._cache.gather(layer_id, self._page_table)
        if self._total_len:
            k = k[:, : self._total_len]
            v = v[:, : self._total_len]
        return k, v
@@ -0,0 +1,40 @@
 """Inference core: cache, executor, scheduler, task management."""
 from astrai.inference.core.cache import (
    Allocator,
    CacheView,
    ContiguousCache,
    ContiguousCacheView,
    KVCache,
    PageCache,
    PageCacheView,
    PagePool,
    PrefixCache,
    Storage,
    TaskTable,
    page_hash,
 )
 from astrai.inference.core.executor import Executor
 from astrai.inference.core.scheduler import InferenceScheduler
 from astrai.inference.core.task import STOP, Task, TaskManager, TaskStatus
 __all__ = [
    "Allocator",
    "CacheView",
    "KVCache",
    "ContiguousCache",
    "ContiguousCacheView",
    "PageCache",
    "PageCacheView",
    "PagePool",
    "PrefixCache",
    "Storage",
    "TaskTable",
    "page_hash",
    "Executor",
    "InferenceScheduler",
    "STOP",
    "Task",
    "TaskManager",
    "TaskStatus",
 ]
@@ -0,0 +1,533 @@
 import threading
 from abc import ABC, abstractmethod
 from collections import OrderedDict
 from typing import Callable, Dict, List, Optional, Tuple
 import torch
 from torch import Tensor
 def page_hash(token_ids: List[int], page_idx: int, page_size: int) -> int:
    start = page_idx * page_size
    end = min(start + page_size, len(token_ids))
    h = 0
    for i in range(start, end):
        h = (h * 31 + token_ids[i]) & 0xFFFFFFFFFFFFFFFF
    return h
 class Allocator:
    """Bitmask-based page allocator with ref-counting and LRU eviction."""
    def __init__(self, n_pages: int):
        self._free_mask = (1 << n_pages) - 1
        self._refs: List[int] = [0] * n_pages
        self._lru: OrderedDict[int, None] = OrderedDict()
        self.on_evict: Optional[Callable[[int], None]] = None
        self._lock = threading.Lock()
    def alloc(self) -> int:
        with self._lock:
            if self._free_mask:
                lsb = self._free_mask & -self._free_mask
                idx = lsb.bit_length() - 1
                self._free_mask ^= lsb
                self._refs[idx] = 1
                return idx
            if self._lru:
                idx, _ = self._lru.popitem(last=False)
                if self.on_evict:
                    self.on_evict(idx)
                self._refs[idx] = 1
                self._free_mask &= ~(1 << idx)
                return idx
            return -1
    def free(self, idx: int, keep_cached: bool = False):
        with self._lock:
            self._refs[idx] -= 1
            if self._refs[idx] == 0:
                if keep_cached:
                    self._lru[idx] = None
                else:
                    self._free_mask |= 1 << idx
    def inc_ref(self, idx: int):
        with self._lock:
            self._refs[idx] += 1
            self._lru.pop(idx, None)
    def ref_count(self, idx: int) -> int:
        with self._lock:
            return self._refs[idx]
    def touch(self, idx: int):
        with self._lock:
            if idx in self._lru:
                self._lru.move_to_end(idx)
 class PrefixCache:
    """Hash-based prefix matching: maps page hashes to physical page indices."""
    def __init__(self, page_size: int):
        self._page_size = page_size
        self._page_to_hash: Dict[int, int] = {}
        self._hash_to_page: Dict[int, int] = {}
        self._lock = threading.Lock()
    def evict(self, idx: int):
        with self._lock:
            h = self._page_to_hash.pop(idx, None)
            if h is not None:
                self._hash_to_page.pop(h, None)
    def has_page(self, idx: int) -> bool:
        with self._lock:
            return idx in self._page_to_hash
    def lookup(self, token_ids: List[int]) -> List[int]:
        with self._lock:
            full_pages = len(token_ids) // self._page_size
            hits: List[int] = []
            for i in range(full_pages):
                h = page_hash(token_ids, i, self._page_size)
                p = self._hash_to_page.get(h)
                if p is None:
                    break
                hits.append(p)
            return hits
    def record(self, page_idx: int, token_ids: List[int], logical_page_idx: int):
        with self._lock:
            h = page_hash(token_ids, logical_page_idx, self._page_size)
            old_h = self._page_to_hash.pop(page_idx, None)
            if old_h is not None:
                self._hash_to_page.pop(old_h, None)
            self._page_to_hash[page_idx] = h
            self._hash_to_page[h] = page_idx
 class PagePool:
    """Orchestrates allocator (page management) and PrefixCache (content addressing)."""
    def __init__(self, allocator: Allocator, prefix: PrefixCache):
        self._alloc = allocator
        self._prefix = prefix
        self._alloc.on_evict = prefix.evict
    @property
    def allocator(self) -> Allocator:
        return self._alloc
    @property
    def prefix(self) -> PrefixCache:
        return self._prefix
    def alloc(self) -> int:
        return self._alloc.alloc()
    def free(self, idx: int):
        keep = self._prefix.has_page(idx)
        self._alloc.free(idx, keep_cached=keep)
        if not keep:
            self._prefix.evict(idx)
    def inc_ref(self, idx: int):
        self._alloc.inc_ref(idx)
    def lookup(self, token_ids: List[int]) -> List[int]:
        hits = self._prefix.lookup(token_ids)
        for p in hits:
            self._alloc.touch(p)
        return hits
    def record(self, page_idx: int, token_ids: List[int], logical_page_idx: int):
        self._prefix.record(page_idx, token_ids, logical_page_idx)
 class TaskTable:
    """Maps task_ids to page tables and cached token counts."""
    def __init__(self, page_size: int):
        self._page_size = page_size
        self._pages: Dict[str, List[int]] = {}
        self._cached: Dict[str, int] = {}
        self._lock = threading.Lock()
    def set(self, task_id: str, page_table: List[int], cached: int):
        with self._lock:
            self._pages[task_id] = page_table
            self._cached[task_id] = cached
    def get(self, task_id: str) -> List[int]:
        with self._lock:
            return self._pages.get(task_id, [])
    def get_cached(self, task_id: str) -> int:
        with self._lock:
            return self._cached.get(task_id, 0)
    def pop(self, task_id: str) -> Tuple[List[int], int]:
        with self._lock:
            pages = self._pages.pop(task_id, [])
            cached = self._cached.pop(task_id, 0)
            return pages, cached
    def get_ref(self, task_id: str) -> List[int]:
        with self._lock:
            return self._pages.setdefault(task_id, [])
    def table_tensor(self, task_ids: List[str], device: torch.device) -> Tensor:
        with self._lock:
            states = [self._pages.get(tid, []) for tid in task_ids]
            max_pages = max((len(s) for s in states), default=0)
            rows = [s + [-1] * (max_pages - len(s)) for s in states]
            return torch.tensor(rows, dtype=torch.long, device=device)
 class Storage:
    """KV-cache tensor storage with paged write/gather."""
    def __init__(
        self,
        n_layers: int,
        n_pages: int,
        page_size: int,
        n_kv_heads: int,
        head_dim: int,
        device: torch.device,
        dtype: torch.dtype,
    ):
        self.page_size = page_size
        self.k_cache = torch.empty(
            (n_layers, n_pages, page_size, n_kv_heads, head_dim),
            device=device,
            dtype=dtype,
        )
        self.v_cache = torch.empty(
            (n_layers, n_pages, page_size, n_kv_heads, head_dim),
            device=device,
            dtype=dtype,
        )
    def write(
        self,
        layer_id: int,
        page_table: Tensor,
        start_pos: int,
        k: Tensor,
        v: Tensor,
    ):
        seq_len = k.size(1)
        if seq_len == 0:
            return
        page_size = self.page_size
        written = 0
        first_page = start_pos // page_size
        last_page = (start_pos + seq_len - 1) // page_size
        for pi in range(first_page, last_page + 1):
            phys_pages = page_table[:, pi]
            page_start = pi * page_size
            write_start = max(page_start, start_pos)
            write_end = min(page_start + page_size, start_pos + seq_len)
            offset = write_start - page_start
            chunk = write_end - write_start
            valid = phys_pages >= 0
            if not valid.all():
                if valid.any():
                    valid_pages = phys_pages[valid]
                    self.k_cache[layer_id, valid_pages, offset : offset + chunk] = k[
                        valid, written : written + chunk
                    ]
                    self.v_cache[layer_id, valid_pages, offset : offset + chunk] = v[
                        valid, written : written + chunk
                    ]
                written += chunk
                continue
            self.k_cache[layer_id, phys_pages, offset : offset + chunk] = k[
                :, written : written + chunk
            ]
            self.v_cache[layer_id, phys_pages, offset : offset + chunk] = v[
                :, written : written + chunk
            ]
            written += chunk
    def gather(
        self, layer_id: int, page_table: Tensor, total_len: int
    ) -> Tuple[Tensor, Tensor]:
        safe = page_table.clamp(min=0)
        k = self.k_cache[layer_id, safe]
        v = self.v_cache[layer_id, safe]
        k = k.flatten(1, 2)
        v = v.flatten(1, 2)
        if (page_table < 0).any():
            invalid = (
                (page_table < 0)
                .unsqueeze(-1)
                .expand(-1, -1, self.page_size)
                .flatten(1, 2)
            )
            invalid = invalid[:, :, None, None].expand_as(k)
            k = k.masked_fill(invalid, 0.0)
            v = v.masked_fill(invalid, 0.0)
        k = k[:, :total_len]
        v = v[:, :total_len]
        return k, v
 class CacheView(ABC):
    """Abstract view passed to attention layers for KV-cache I/O."""
    @abstractmethod
    def write(self, layer_id: int, k: Tensor, v: Tensor): ...
    @abstractmethod
    def gather(self, layer_id: int) -> Tuple[Tensor, Tensor]: ...
 class KVCache(ABC):
    """Abstract KV-cache facade for scheduler/executor."""
    @abstractmethod
    def task_alloc(self, task_id: str, prompt_ids: List[int]) -> bool: ...
    @abstractmethod
    def task_free(self, task_id: str): ...
    @abstractmethod
    def task_extend(self, task_id: str, pos: int) -> bool: ...
    @abstractmethod
    def bind_tasks(
        self,
        task_ids: List[str],
        total_len: int,
        device: torch.device,
        write_positions: Optional[Tensor] = None,
    ) -> CacheView: ...
    def task_cached(self, task_id: str) -> int:
        return 0
    def task_record_hashes(
        self, task_id: str, prompt_ids: List[int], start_logical_page: int = 0
    ): ...
 class PageCacheView(CacheView):
    """Bundles Storage + page_table + total_len for attention layers."""
    def __init__(self, storage: Storage, page_table: Tensor, total_len: int = 0):
        self._storage = storage
        self._page_table = page_table
        self._total_len = total_len
    def write(self, layer_id: int, k: Tensor, v: Tensor):
        start_pos = self._total_len - k.size(1)
        self._storage.write(layer_id, self._page_table, start_pos, k, v)
    def gather(self, layer_id: int) -> Tuple[Tensor, Tensor]:
        return self._storage.gather(layer_id, self._page_table, self._total_len)
 class PageCache(KVCache):
    """Paged KV-cache with prefix sharing."""
    def __init__(
        self,
        n_layers: int,
        n_pages: int,
        page_size: int,
        n_kv_heads: int,
        head_dim: int,
        device: torch.device,
        dtype: torch.dtype,
    ):
        self.page_size = page_size
        self._pool = PagePool(Allocator(n_pages), PrefixCache(page_size))
        self._table = TaskTable(page_size)
        self._storage = Storage(
            n_layers, n_pages, page_size, n_kv_heads, head_dim, device, dtype
        )
    def task_alloc(self, task_id: str, prompt_ids: List[int]) -> bool:
        hits = self._pool.lookup(prompt_ids)
        cached = len(hits) * self.page_size
        for p in hits:
            self._pool.inc_ref(p)
        remaining = len(prompt_ids) - cached
        n_new = (
            (remaining + self.page_size - 1) // self.page_size if remaining > 0 else 0
        )
        new_pages: List[int] = []
        if n_new > 0:
            for _ in range(n_new):
                p = self._pool.alloc()
                if p < 0:
                    for hp in hits:
                        self._pool.free(hp)
                    for np in new_pages:
                        self._pool.free(np)
                    return False
                new_pages.append(p)
        self._table.set(task_id, hits + new_pages, cached)
        return True
    def task_free(self, task_id: str):
        page_table, _ = self._table.pop(task_id)
        for idx in page_table:
            self._pool.free(idx)
    def task_extend(self, task_id: str, pos: int) -> bool:
        page_table = self._table.get(task_id)
        needed = (pos + 1 + self.page_size - 1) // self.page_size
        while len(page_table) < needed:
            p = self._pool.alloc()
            if p < 0:
                return False
            page_table.append(p)
        return True
    def task_cached(self, task_id: str) -> int:
        return self._table.get_cached(task_id)
    def task_record_hashes(
        self, task_id: str, prompt_ids: List[int], start_logical_page: int = 0
    ):
        page_table = self._table.get(task_id)
        full_pages = len(prompt_ids) // self.page_size
        for i in range(start_logical_page, full_pages):
            self._pool.record(page_table[i], prompt_ids, i)
    def bind_tasks(
        self,
        task_ids: List[str],
        total_len: int,
        device: torch.device,
        write_positions: Optional[Tensor] = None,
    ) -> PageCacheView:
        page_table = self._table.table_tensor(task_ids, device)
        return PageCacheView(self._storage, page_table, total_len)
 class ContiguousCacheView(CacheView):
    """Contiguous KV-cache view for attention layers."""
    def __init__(
        self,
        cache: "ContiguousCache",
        batch_indices: Tensor,
        total_len: int = 0,
        write_positions: Optional[Tensor] = None,
    ):
        self._cache = cache
        self._batch_indices = batch_indices
        self._total_len = total_len
        self._write_positions = write_positions
    def write(self, layer_id: int, k: Tensor, v: Tensor):
        seq_len = k.size(1)
        indices = self._batch_indices
        if self._write_positions is not None and seq_len == 1:
            pos = self._write_positions
            self._cache.k[layer_id, indices, pos] = k.squeeze(1)
            self._cache.v[layer_id, indices, pos] = v.squeeze(1)
            for s, p in zip(indices.tolist(), pos.tolist()):
                cur = self._cache._slot_len.get(s, 0)
                if p + 1 > cur:
                    self._cache._slot_len[s] = p + 1
        else:
            start_pos = self._total_len - seq_len
            self._cache.k[layer_id, indices, start_pos : start_pos + seq_len] = k
            self._cache.v[layer_id, indices, start_pos : start_pos + seq_len] = v
            new_len = start_pos + seq_len
            for s in indices.tolist():
                cur = self._cache._slot_len.get(s, 0)
                if new_len > cur:
                    self._cache._slot_len[s] = new_len
    def gather(self, layer_id: int) -> Tuple[Tensor, Tensor]:
        max_len = max(
            self._cache._slot_len.get(int(s), 0) for s in self._batch_indices.tolist()
        )
        indices = self._batch_indices
        k = self._cache.k[layer_id, indices, :max_len]
        v = self._cache.v[layer_id, indices, :max_len]
        return k, v
 class ContiguousCache(KVCache):
    """Contiguous per-slot KV cache (default implementation)."""
    def __init__(
        self,
        n_layers: int,
        max_batch_size: int,
        max_seq_len: int,
        n_kv_heads: int,
        head_dim: int,
        device: torch.device,
        dtype: torch.dtype,
    ):
        self.max_seq_len = max_seq_len
        self.k = torch.zeros(
            n_layers,
            max_batch_size,
            max_seq_len,
            n_kv_heads,
            head_dim,
            device=device,
            dtype=dtype,
        )
        self.v = torch.zeros(
            n_layers,
            max_batch_size,
            max_seq_len,
            n_kv_heads,
            head_dim,
            device=device,
            dtype=dtype,
        )
        self._slot_len: Dict[int, int] = {}
        self._task_slot: Dict[str, int] = {}
        self._free_slots = list(range(max_batch_size))
        self._device = device
    def task_alloc(self, task_id: str, prompt_ids: List[int]) -> bool:
        if not self._free_slots:
            return False
        slot = self._free_slots.pop(0)
        self._task_slot[task_id] = slot
        self._slot_len[slot] = 0
        return True
    def task_free(self, task_id: str):
        slot = self._task_slot.pop(task_id, None)
        if slot is not None:
            self._slot_len.pop(slot, None)
            self._free_slots.append(slot)
    def task_extend(self, task_id: str, pos: int) -> bool:
        return pos < self.max_seq_len
    def task_cached(self, task_id: str) -> int:
        slot = self._task_slot.get(task_id)
        if slot is None:
            return 0
        return self._slot_len.get(slot, 0)
    def bind_tasks(
        self,
        task_ids: List[str],
        total_len: int,
        device: torch.device,
        write_positions: Optional[Tensor] = None,
    ) -> ContiguousCacheView:
        slots = [self._task_slot[tid] for tid in task_ids]
        batch_indices = torch.tensor(slots, dtype=torch.long, device=device)
        return ContiguousCacheView(
            self, batch_indices, total_len, write_positions=write_positions
        )
@@ -0,0 +1,127 @@
 import logging
 from typing import List, Optional
 import torch
 from astrai.inference.core.cache import KVCache
 from astrai.inference.core.task import Task
 from astrai.inference.sample import sample
 from astrai.model.automodel import AutoModel
 from astrai.tokenize.tokenizer import AutoTokenizer
 logger = logging.getLogger(__name__)
 class Executor:
    """Model forward passes for prefill and decode phases."""
    def __init__(
        self,
        model: AutoModel,
        tokenizer: AutoTokenizer,
        kv_cache: KVCache,
        device: Optional[str] = None,
        dtype: Optional[torch.dtype] = None,
    ):
        self.model = model
        self.tokenizer = tokenizer
        self.kv_cache = kv_cache
        self.device = device or next(model.parameters()).device
        self.dtype = dtype or next(model.parameters()).dtype
    def execute_prefill(self, tasks: List[Task], prompt_len: int, start_pos: int = 0):
        if start_pos >= prompt_len:
            return
        tasks = sorted(tasks, key=lambda t: t.task_id)
        batch_sz = len(tasks)
        input_ids = torch.tensor(
            [t.prompt_ids[start_pos:prompt_len] for t in tasks],
            dtype=torch.long,
            device=self.device,
        )
        task_ids = [t.task_id for t in tasks]
        with torch.inference_mode():
            self.model(
                input_ids,
                position_ids=torch.arange(
                    start_pos, prompt_len, dtype=torch.long, device=self.device
                )
                .unsqueeze(0)
                .expand(batch_sz, -1),
                paged_cache=self.kv_cache.bind_tasks(task_ids, prompt_len, self.device),
            )
    def execute_decode(self, tasks: List[Task]) -> List[int]:
        if not tasks:
            return []
        input_ids = torch.tensor(
            [t.output_ids[-1] if t.output_ids else t.prompt_ids[-1] for t in tasks],
            dtype=torch.long,
            device=self.device,
        )
        position_ids = torch.tensor(
            [t.next_pos for t in tasks], dtype=torch.long, device=self.device
        )
        total_len = position_ids.max().item() + 1
        task_ids = [t.task_id for t in tasks]
        temperatures = torch.tensor([t.temperature for t in tasks], device=self.device)
        top_ks = torch.tensor([t.top_k for t in tasks], device=self.device)
        top_ps = torch.tensor([t.top_p for t in tasks], device=self.device)
        freq_penalties = torch.tensor(
            [t.frequency_penalty for t in tasks], device=self.device
        )
        history_lists = []
        mask_lists = []
        for t in tasks:
            window = t.rep_window
            prompt_part = t.prompt_ids[-window:]
            ids = prompt_part + t.output_ids
            history_lists.append(ids)
            mask_lists.append([True] * len(ids))
        max_len = max(len(h) for h in history_lists)
        padded_ids = torch.zeros(
            len(tasks), max_len, dtype=torch.long, device=self.device
        )
        padded_mask = torch.zeros(
            len(tasks), max_len, dtype=torch.bool, device=self.device
        )
        for i, (h, m) in enumerate(zip(history_lists, mask_lists)):
            padded_ids[i, : len(h)] = torch.tensor(
                h, dtype=torch.long, device=self.device
            )
            padded_mask[i, : len(m)] = torch.tensor(
                m, dtype=torch.bool, device=self.device
            )
        with torch.inference_mode():
            outputs = self.model(
                input_ids.unsqueeze(1),
                paged_cache=self.kv_cache.bind_tasks(
                    task_ids,
                    total_len,
                    self.device,
                    write_positions=position_ids,
                ),
                position_ids=position_ids.unsqueeze(1),
            )
            logits = outputs["logits"][:, -1, :]
        return sample(
            logits,
            temperature=temperatures,
            top_k=top_ks,
            top_p=top_ps,
            frequency_penalty=freq_penalties,
            input_ids=padded_ids,
            input_mask=padded_mask,
        ).tolist()
@@ -0,0 +1,199 @@
 import logging
 import threading
 from typing import Any, Dict, List, Optional, Tuple
 import torch
 from astrai.inference.core.cache import ContiguousCache, KVCache
 from astrai.inference.core.executor import Executor
 from astrai.inference.core.task import STOP, Task, TaskManager, TaskStatus
 from astrai.model.automodel import AutoModel
 from astrai.tokenize.tokenizer import AutoTokenizer
 logger = logging.getLogger(__name__)
 class InferenceScheduler:
    """Continuous batching loop: cleanup -> refill -> prefill -> decode (all groups)."""
    def __init__(
        self,
        model: AutoModel,
        tokenizer: AutoTokenizer,
        max_batch_size: int = 16,
        max_seq_len: Optional[int] = None,
        max_prompt_len: int = 2048,
        device: Optional[str] = None,
        dtype: Optional[torch.dtype] = None,
        cache: Optional[KVCache] = None,
    ):
        config = model.config
        if max_seq_len is not None:
            self.max_seq_len = max_seq_len
        elif config.max_len is not None:
            self.max_seq_len = config.max_len
        else:
            raise ValueError(
                "max_seq_len must be provided either as argument "
                "or in model config (config.max_len)"
            )
        self.device = device or next(model.parameters()).device
        self.dtype = dtype or next(model.parameters()).dtype
        head_dim = config.dim // config.n_heads
        if cache is not None:
            self._cache = cache
        else:
            self._cache = ContiguousCache(
                config.n_layers,
                max_batch_size,
                self.max_seq_len,
                config.n_kv_heads,
                head_dim,
                self.device,
                self.dtype,
            )
        self._task_mgr = TaskManager(
            tokenizer=tokenizer,
            max_batch_size=max_batch_size,
            max_seq_len=self.max_seq_len,
            max_prompt_len=max_prompt_len,
        )
        self._executor = Executor(
            model=model,
            tokenizer=tokenizer,
            kv_cache=self._cache,
            device=self.device,
            dtype=self.dtype,
        )
        self._stop_event = threading.Event()
        self._loop_thread: Optional[threading.Thread] = None
    def add_task(self, prompt: str, **kwargs) -> str:
        return self._task_mgr.add_task(prompt, **kwargs)
    def remove_task(self, task_id: str):
        for task in self._task_mgr.remove_task(task_id):
            self._cache.task_free(task.task_id)
    def get_stats(self) -> Dict[str, Any]:
        return self._task_mgr.get_stats()
    def _run_generation_loop(self):
        stop_ids = self._task_mgr.tokenizer.stop_ids
        cache = self._cache
        try:
            while not self._stop_event.is_set():
                finished = self._task_mgr.remove_finished_tasks(stop_ids)
                for task in finished:
                    cache.task_free(task.task_id)
                active = self._task_mgr.get_active_tasks()
                available = self._task_mgr.max_batch_size - len(active)
                if available > 0:
                    candidates = self._task_mgr.pull_candidates(available)
                    failed = []
                    for task in candidates:
                        if cache.task_alloc(task.task_id, task.prompt_ids):
                            self._task_mgr.activate(task)
                        else:
                            failed.append(task)
                    if failed:
                        self._task_mgr.return_to_waiting(failed)
                if not self._task_mgr.has_work():
                    self._task_mgr.wait_for_tasks(timeout=1.0)
                    continue
                to_prefill = [
                    t
                    for t in self._task_mgr.get_active_tasks()
                    if t.output_tokens == 0
                    and cache.task_cached(t.task_id) < len(t.prompt_ids)
                ]
                if to_prefill:
                    for t in to_prefill:
                        t.input_tokens = len(t.prompt_ids)
                    groups: Dict[Tuple[int, int], List[Task]] = {}
                    for t in to_prefill:
                        key = (
                            len(t.prompt_ids),
                            cache.task_cached(t.task_id),
                        )
                        groups.setdefault(key, []).append(t)
                    for (prompt_len, start_pos), group in groups.items():
                        self._executor.execute_prefill(group, prompt_len, start_pos)
                        start_logical_page = start_pos // getattr(
                            cache, "page_size", 64
                        )
                        for t in group:
                            cache.task_record_hashes(
                                t.task_id, t.prompt_ids, start_logical_page
                            )
                decode_tasks = self._task_mgr.get_active_tasks()
                valid: List[Task] = []
                for t in sorted(decode_tasks, key=lambda t: t.task_id):
                    if cache.task_extend(t.task_id, t.next_pos):
                        valid.append(t)
                    else:
                        t.status = TaskStatus.ABORTED
                        self._task_mgr.invoke_callback(t.task_id, STOP)
                if valid:
                    next_tokens = self._executor.execute_decode(valid)
                    for t, ntok in zip(valid, next_tokens):
                        t.output_ids.append(ntok)
                        t.output_tokens += 1
                        new_text = t.decode_new_token(self._task_mgr.tokenizer)
                        if new_text:
                            self._task_mgr.invoke_callback(t.task_id, new_text)
                    for t in valid:
                        if t.is_finished(stop_ids):
                            remaining = t.flush_remaining(self._task_mgr.tokenizer)
                            if remaining:
                                self._task_mgr.invoke_callback(t.task_id, remaining)
                            self._task_mgr.invoke_callback(t.task_id, STOP)
        except Exception as e:
            self._stop_event.set()
            logger.error(f"Scheduler loop crashed: {e}", exc_info=True)
            for task in self._task_mgr.get_active_tasks():
                self._task_mgr.invoke_callback(task.task_id, STOP)
                cache.task_free(task.task_id)
            for task in self._task_mgr.get_waiting_tasks():
                self._task_mgr.invoke_callback(task.task_id, STOP)
            self._task_mgr.clear_queues()
    def start(self):
        if self._loop_thread is not None and self._loop_thread.is_alive():
            return
        self._stop_event.clear()
        t = threading.Thread(target=self._run_generation_loop, daemon=True)
        t.start()
        self._loop_thread = t
    def stop(self):
        self._stop_event.set()
        self._task_mgr.wake()
        if self._loop_thread is not None:
            self._loop_thread.join(timeout=2.0)
        self._loop_thread = None
        for task in self._task_mgr.get_active_tasks():
            self._task_mgr.invoke_callback(task.task_id, STOP)
            self._cache.task_free(task.task_id)
        for task in self._task_mgr.get_waiting_tasks():
            self._task_mgr.invoke_callback(task.task_id, STOP)
        self._task_mgr.clear_queues()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
@@ -0,0 +1,286 @@
 import logging
 import threading
 import time
 import uuid
 from collections import deque
 from enum import Enum
 from typing import Any, Callable, Deque, Dict, List, Optional
 from astrai.tokenize.tokenizer import AutoTokenizer
 logger = logging.getLogger(__name__)
 STOP = object()
 class StreamDecoder:
    """Incremental decoder for byte-level BPE streaming.
    Byte-level BPE may split a single Unicode character (e.g. em-dash,
    smart quotes) across multiple tokens. Decoding such a token in
    isolation produces U+FFFD (replacement char). This decoder
    accumulates token IDs and only emits text once the trailing
    characters are complete, buffering incomplete multi-byte sequences
    until the next token arrives.
    """
    __slots__ = ("_tokenizer", "_ids", "_emitted")
    def __init__(self, tokenizer: AutoTokenizer):
        self._tokenizer = tokenizer
        self._ids: List[int] = []
        self._emitted: str = ""
    def push(self, token_id: int) -> str:
        """Append a token ID and return newly completed text.
        Returns "" while a multi-byte character is still incomplete.
        """
        self._ids.append(token_id)
        full = self._tokenizer.decode(self._ids, skip_special_tokens=True)
        if full.endswith("\ufffd"):
            return ""
        if len(full) > len(self._emitted):
            diff = full[len(self._emitted) :]
            self._emitted = full
            return diff
        return ""
 class TaskStatus(Enum):
    """Task lifecycle states."""
    PENDING = "pending"
    RUNNING = "running"
    FINISHED = "finished"
    ABORTED = "aborted"
 class Task:
    """Single generation request: prompt, sampling params, output state."""
    def __init__(
        self,
        task_id: str,
        prompt_ids: List[int],
        max_tokens: Optional[int] = None,
        temperature: float = 1.0,
        top_p: float = 1.0,
        top_k: int = 50,
        frequency_penalty: float = 0.0,
        rep_window: int = 64,
    ):
        self.task_id = task_id
        self.prompt_ids = prompt_ids
        self.max_tokens = max_tokens
        self.temperature = temperature
        self.top_p = top_p
        self.top_k = top_k
        self.frequency_penalty = frequency_penalty
        self.rep_window = rep_window
        self.status = TaskStatus.PENDING
        self.output_ids: List[int] = []
        self.input_tokens: int = 0
        self.output_tokens: int = 0
        self.arrival_time = time.time()
        self.finish_time: Optional[float] = None
        self._decoder: Optional[StreamDecoder] = None
    def decode_new_token(self, tokenizer: AutoTokenizer) -> str:
        """Decode the last appended output token, buffering incomplete
        multi-byte sequences across calls.
        Lazily creates a :class:`StreamDecoder` on first use.
        """
        if self._decoder is None:
            self._decoder = StreamDecoder(tokenizer)
        return self._decoder.push(self.output_ids[-1])
    def flush_remaining(self, tokenizer: AutoTokenizer) -> str:
        """Emit any text still buffered in the decoder.
        Called when generation terminates (max_tokens reached, stop
        sequence, or external removal) to avoid dropping a final
        incomplete-looking fragment that is actually complete when
        adjacent to the stop token.
        """
        if self._decoder is None or not self.output_ids:
            return ""
        full = tokenizer.decode(self.output_ids, skip_special_tokens=True)
        if len(full) > len(self._decoder._emitted):
            diff = full[len(self._decoder._emitted) :]
            self._decoder._emitted = full
            return diff
        return ""
    @property
    def next_pos(self) -> int:
        return self.input_tokens + len(self.output_ids)
    def is_finished(self, stop_ids: List[int]) -> bool:
        if self.max_tokens is not None and self.output_tokens >= self.max_tokens:
            return True
        if self.output_ids and self.output_ids[-1] in stop_ids:
            return True
        return False
 class TaskManager:
    """Thread-safe task queues and lifecycle transitions (no page ops)."""
    def __init__(
        self,
        tokenizer: AutoTokenizer,
        max_batch_size: int = 16,
        max_seq_len: int = 8192,
        max_prompt_len: int = 512,
    ):
        self.tokenizer = tokenizer
        self.max_batch_size = max_batch_size
        self.max_seq_len = max_seq_len
        self.max_prompt_len = max_prompt_len
        self.waiting_queue: Deque[Task] = deque()
        self.active_tasks: List[Task] = []
        self._callbacks: Dict[str, Callable[[str], None]] = {}
        self._task_event = threading.Event()
        self._lock = threading.Lock()
        self._total_tasks = 0
        self._total_tokens = 0
    def add_task(
        self,
        prompt: str,
        max_tokens: Optional[int] = None,
        temperature: float = 1.0,
        top_p: float = 1.0,
        top_k: int = 50,
        frequency_penalty: float = 0.0,
        rep_window: int = 64,
        stream_callback: Optional[Callable[[str], None]] = None,
    ) -> str:
        task_id = f"task_{int(time.time())}_{uuid.uuid4().hex[:8]}"
        prompt_ids = self.tokenizer.encode(prompt)
        if len(prompt_ids) > self.max_prompt_len:
            prompt_ids = prompt_ids[-self.max_prompt_len :]
        if len(prompt_ids) >= self.max_seq_len:
            if stream_callback:
                stream_callback(STOP)
            return task_id
        if max_tokens is None:
            max_tokens = self.max_seq_len - len(prompt_ids)
        else:
            max_tokens = min(max_tokens, self.max_seq_len - len(prompt_ids))
        task = Task(
            task_id=task_id,
            prompt_ids=prompt_ids,
            max_tokens=max_tokens,
            temperature=temperature,
            top_p=top_p,
            top_k=top_k,
            frequency_penalty=frequency_penalty,
            rep_window=rep_window,
        )
        with self._lock:
            self.waiting_queue.append(task)
            self._total_tasks += 1
            if stream_callback:
                self._callbacks[task_id] = stream_callback
        self._task_event.set()
        return task_id
    def remove_task(self, task_id: str) -> List[Task]:
        with self._lock:
            removed_active = [t for t in self.active_tasks if t.task_id == task_id]
            self.waiting_queue = deque(
                t for t in self.waiting_queue if t.task_id != task_id
            )
            self.active_tasks = [t for t in self.active_tasks if t.task_id != task_id]
            self._callbacks.pop(task_id, None)
        return removed_active
    def invoke_callback(self, task_id: str, token: str):
        cb = self._callbacks.get(task_id)
        if cb:
            cb(token)
    def get_stats(self) -> Dict[str, Any]:
        return {
            "total_tasks": self._total_tasks,
            "total_tokens": self._total_tokens,
            "active_tasks": len(self.active_tasks),
            "waiting_queue": len(self.waiting_queue),
        }
    def remove_finished_tasks(self, stop_ids: List[int]) -> List[Task]:
        with self._lock:
            finished = []
            for task in self.active_tasks:
                if task.status == TaskStatus.ABORTED:
                    task.finish_time = time.time()
                    finished.append(task)
                elif task.is_finished(stop_ids):
                    task.status = TaskStatus.FINISHED
                    task.finish_time = time.time()
                    finished.append(task)
                    self._total_tokens += task.output_tokens
            self.active_tasks = [
                t
                for t in self.active_tasks
                if t.status not in (TaskStatus.FINISHED, TaskStatus.ABORTED)
            ]
            return finished
    def pull_candidates(self, n: int) -> List[Task]:
        to_add: List[Task] = []
        with self._lock:
            take = min(n, len(self.waiting_queue))
            for _ in range(take):
                to_add.append(self.waiting_queue.popleft())
        return to_add
    def activate(self, task: Task):
        task.status = TaskStatus.RUNNING
        with self._lock:
            self.active_tasks.append(task)
    def return_to_waiting(self, tasks: List[Task]):
        with self._lock:
            for task in reversed(tasks):
                self.waiting_queue.appendleft(task)
    def has_work(self) -> bool:
        return bool(self.active_tasks or self.waiting_queue)
    def wait_for_tasks(self, timeout: float = 1.0):
        with self._lock:
            if self.waiting_queue or self.active_tasks:
                return
            self._task_event.clear()
        self._task_event.wait(timeout=timeout)
    def get_active_tasks(self) -> List[Task]:
        with self._lock:
            return list(self.active_tasks)
    def get_waiting_tasks(self) -> List[Task]:
        with self._lock:
            return list(self.waiting_queue)
    def clear_queues(self):
        with self._lock:
            self.waiting_queue.clear()
            self.active_tasks.clear()
            self._callbacks.clear()
    def wake(self):
        self._task_event.set()
@@ -1,132 +1,34 @@
-"""Unified inference engine for continuous batching.
+"""Unified inference engine for continuous batching."""
 Layers:
  - GenerationParams:    Immutable value object for sampling parameters.
  - GenerationRequest:   User-facing request DTO with validation.
  - _Result:             Thread-safe token accumulator (Observer pattern).
  - InferenceEngine:     Facade over InferenceScheduler + async wrapper.
 """
 import asyncio
 import gc
 import threading
-from dataclasses import dataclass
+from typing import Any, AsyncGenerator, Dict, Generator, List, Optional, Tuple, Union
 from typing import Any, AsyncGenerator, Dict, Generator, List, Optional, Union
 import torch
 import torch.nn as nn
-from astrai.inference.cache import STOP
+from astrai.inference.core.cache import KVCache
-from astrai.inference.scheduler import InferenceScheduler
+from astrai.inference.core.scheduler import InferenceScheduler
 from astrai.inference.core.task import STOP
 from astrai.tokenize import AutoTokenizer
-@dataclass(frozen=True)
+class GenerateResult:
-class GenerationParams:
+    """Thread-safe token accumulator for streaming and non-streaming modes."""
    """Immutable value object for sampling hyperparameters."""
    top_k: int = 50
    top_p: float = 1.0
    temperature: float = 1.0
    max_tokens: int = 1024
 class GenerationRequest:
    """Request parameters for text generation.
    Encapsulates messages, sampling parameters (via GenerationParams),
    and streaming preference for a single generation request.
    """
    def __init__(
        self,
        messages: List[Dict[str, str]],
        top_k: int = 50,
        top_p: float = 1.0,
        temperature: float = 1.0,
        max_len: int = 1024,
        stream: bool = False,
    ):
        """Initializes a generation request.
        Args:
            messages: Conversation history as list of {"role": ..., "content": ...}.
            top_k: Top-k sampling count (0 disables).
            top_p: Nucleus sampling probability threshold.
            temperature: Sampling temperature.
            max_len: Maximum tokens to generate.
            stream: Whether to return output as a token stream.
        """
        self.messages = messages
        self.params = GenerationParams(
            top_k=top_k,
            top_p=top_p,
            temperature=temperature,
            max_tokens=max_len,
        )
        self.stream = stream
        self._validate()
    @property
    def top_k(self) -> int:
        return self.params.top_k
    @property
    def top_p(self) -> float:
        return self.params.top_p
    @property
    def temperature(self) -> float:
        return self.params.temperature
    @property
    def max_len(self) -> int:
        return self.params.max_tokens
    def _validate(self):
        """Validates sampling parameter ranges."""
        if not (isinstance(self.top_k, int) and self.top_k >= 0):
            raise ValueError("top_k must be a non-negative integer")
        if not (0.0 <= self.top_p <= 1.0):
            raise ValueError("top_p must be a float between 0.0 and 1.0")
        if not (isinstance(self.temperature, (int, float)) and self.temperature >= 0):
            raise ValueError("temperature must be a non-negative number")
 class _Result:
    """Thread-safe token accumulator for streaming and non-streaming modes.
    Supports multiple concurrent generation tasks with per-index result tracking.
    Uses a threading.Condition for efficient completion notification
    and a threading.Event for streaming wakeup.
    """
    def __init__(self, count: int = 1):
        """Initializes the accumulator.
        Args:
            count: Number of concurrent generation tasks to track.
        """
        self._cond = threading.Condition()
        self._event = threading.Event()
-        self.tokens: List[str] = []
+        self.tokens: List[Tuple[int, str]] = []
        self.results: List[str] = [""] * count
        self._done: List[bool] = [False] * count
        self._completed = 0
        self._total = count
    def append(self, token: str, idx: int = 0):
        """Appends a token to the result buffer.
        In non-streaming mode, tokens are concatenated into results[idx].
        The sentinel STOP marks a task as complete.
        Args:
            token: The decoded token string, or STOP sentinel.
            idx: Index of the generation task this token belongs to.
        """
        with self._cond:
-            self.tokens.append(token)
+            self.tokens.append((idx, token))
            if token is not STOP:
                self.results[idx] += token
            else:
@@ -136,12 +38,7 @@ class _Result:
                    self._cond.notify_all()
            self._event.set()
-    def pop_all(self) -> List[str]:
+    def pop_all(self) -> List[Tuple[int, str]]:
        """Returns and clears all accumulated tokens.
        Returns:
            List of token strings since the last call.
        """
        with self._cond:
            out = self.tokens.copy()
            self.tokens.clear()
@@ -150,45 +47,63 @@ class _Result:
            return out
    def wait(self, timeout: Optional[float] = None) -> bool:
        """Blocks until new tokens arrive or the timeout expires.
        Args:
            timeout: Maximum wait time in seconds (None = infinite).
        Returns:
            True if the event was set (new data available), False on timeout.
        """
        return self._event.wait(timeout=timeout)
-    def wait_completion(self) -> None:
+    def wait_completion(self, timeout: float = 300.0):
        """Blocks until all tasks complete (non-streaming).
        Uses a Condition to sleep efficiently instead of busy-waiting.
        The calling thread is parked until a STOP signal arrives.
        """
        with self._cond:
-            self._cond.wait_for(lambda: self._completed >= self._total)
+            if not self._cond.wait_for(
                lambda: self._completed >= self._total, timeout=timeout
            ):
                raise TimeoutError(
                    f"Generation timeout after {timeout}s "
                    f"({self._completed}/{self._total} completed)"
                )
    def get_results(self) -> List[str]:
        """Returns all accumulated results for non-streaming mode.
        Returns:
            List of complete generated strings, one per task index.
        """
        with self._cond:
            return self.results.copy()
 class GenerationRequest:
    """Request parameters for text generation."""
    def __init__(
        self,
        messages: List[Dict[str, str]],
        top_k: int = 50,
        top_p: float = 1.0,
        temperature: float = 1.0,
        max_tokens: Optional[int] = None,
        frequency_penalty: float = 0.0,
        rep_window: int = 64,
        stream: bool = False,
    ):
        if not (isinstance(top_k, int) and top_k >= 0):
            raise ValueError("top_k must be a non-negative integer")
        if not (0.0 <= top_p <= 1.0):
            raise ValueError("top_p must be a float between 0.0 and 1.0")
        if not (isinstance(temperature, (int, float)) and temperature >= 0):
            raise ValueError("temperature must be a non-negative number")
        if not (
            isinstance(frequency_penalty, (int, float))
            and -2.0 <= frequency_penalty <= 2.0
        ):
            raise ValueError("frequency_penalty must be between -2.0 and 2.0")
        if not (isinstance(rep_window, int) and rep_window > 0):
            raise ValueError("rep_window must be a positive integer")
        self.messages = messages
        self.top_k = top_k
        self.top_p = top_p
        self.temperature = temperature
        self.max_tokens = max_tokens
        self.frequency_penalty = frequency_penalty
        self.rep_window = rep_window
        self.stream = stream
 class InferenceEngine:
-    """Unified inference engine backed by continuous-batching scheduler.
+    """Unified inference engine backed by continuous-batching scheduler."""
    Usage:
        with InferenceEngine(model, tokenizer) as engine:
            for token in engine.generate("hello", stream=True):
                print(token, end="")
            text = engine.generate("hello")
    """
    def __init__(
        self,
@@ -198,18 +113,8 @@ class InferenceEngine:
        max_seq_len: Optional[int] = None,
        max_prompt_len: int = 2048,
        page_size: int = 128,
        cache: Optional[KVCache] = None,
    ):
        """Initializes the inference engine.
        Args:
            model: The model instance.
            tokenizer: The tokenizer instance.
            max_batch_size: Maximum number of concurrent tasks.
            max_seq_len: Maximum sequence length.
            max_prompt_len: Maximum prompt tokens.
            compile: Whether to compile the model with torch.compile.
            page_size: Number of tokens per KV cache page.
        """
        self.model = model
        self.tokenizer = tokenizer
        self.scheduler = InferenceScheduler(
@@ -218,7 +123,7 @@ class InferenceEngine:
            max_batch_size=max_batch_size,
            max_seq_len=max_seq_len,
            max_prompt_len=max_prompt_len,
-            page_size=page_size,
+            cache=cache,
        )
        self.scheduler.start()
@@ -234,62 +139,58 @@ class InferenceEngine:
        self,
        prompt: Union[str, List[str]],
        stream: bool = False,
-        max_tokens: int = 1024,
+        max_tokens: Optional[int] = None,
        temperature: float = 1.0,
        top_p: float = 1.0,
        top_k: int = 50,
-    ) -> Union[Generator[str, None, None], str, List[str]]:
+        frequency_penalty: float = 0.0,
-        """Generates text from a prompt.
+        rep_window: int = 64,
-
+    ) -> Union[Generator, str, List[str]]:
        Args:
            prompt: Single string or list of strings for batch generation.
            stream: If True, returns a generator yielding tokens one by one.
            max_tokens: Maximum number of tokens to generate.
            temperature: Sampling temperature.
            top_p: Nucleus sampling probability threshold.
            top_k: Top-k sampling count (0 disables).
        Returns:
            Generator (stream=True), single string (non-stream, single prompt),
            or list of strings (non-stream, batch prompts).
        """
        is_batch = isinstance(prompt, list)
        prompts = prompt if is_batch else [prompt]
        if stream:
            return self._generate_streaming(
-                prompts, is_batch, max_tokens, temperature, top_p, top_k
+                prompts,
                is_batch,
                max_tokens,
                temperature,
                top_p,
                top_k,
                frequency_penalty,
                rep_window,
            )
        else:
            return self._generate_non_streaming(
-                prompts, is_batch, max_tokens, temperature, top_p, top_k
+                prompts,
                is_batch,
                max_tokens,
                temperature,
                top_p,
                top_k,
                frequency_penalty,
                rep_window,
            )
    def generate_async(
        self,
        prompt: str,
-        max_tokens: int = 1024,
+        max_tokens: Optional[int] = None,
        temperature: float = 1.0,
        top_p: float = 1.0,
        top_k: int = 50,
        frequency_penalty: float = 0.0,
        rep_window: int = 64,
    ) -> AsyncGenerator[str, None]:
        """Async streaming generator that does not block the event loop.
        Runs the synchronous generator in a background thread pool executor,
        yielding tokens to the async consumer as they arrive.
        Args:
            prompt: Input text to generate from.
            max_tokens: Maximum tokens to generate.
            temperature: Sampling temperature.
            top_p: Nucleus sampling threshold.
            top_k: Top-k sampling count.
        Yields:
            Decoded token strings as they are generated.
        """
        sync_gen = self._generate_streaming(
-            [prompt], False, max_tokens, temperature, top_p, top_k
+            [prompt],
            False,
            max_tokens,
            temperature,
            top_p,
            top_k,
            frequency_penalty,
            rep_window,
        )
        async def _agen():
@@ -304,14 +205,6 @@ class InferenceEngine:
    @staticmethod
    def _next_token(gen: Generator) -> Optional[str]:
        """Retrieves the next token from a synchronous generator.
        Args:
            gen: A synchronous generator yielding token strings.
        Returns:
            The next token, or None if the generator is exhausted.
        """
        try:
            return next(gen)
        except StopIteration:
@@ -320,77 +213,94 @@ class InferenceEngine:
    def generate_with_request(
        self, request: GenerationRequest
    ) -> Union[Generator[str, None, None], str, List[str]]:
        """Generates text from a structured GenerationRequest.
        Applies the chat template to the request's messages before generation.
        Args:
            request: A GenerationRequest with messages and parameters.
        Returns:
            Generator, string, or list of strings (see generate()).
        """
        prompt = self.tokenizer.apply_chat_template(request.messages, tokenize=False)
        return self.generate(
            prompt=prompt,
            stream=request.stream,
-            max_tokens=request.params.max_tokens,
+            max_tokens=request.max_tokens,
-            temperature=request.params.temperature,
+            temperature=request.temperature,
-            top_p=request.params.top_p,
+            top_p=request.top_p,
-            top_k=request.params.top_k,
+            top_k=request.top_k,
            frequency_penalty=request.frequency_penalty,
            rep_window=request.rep_window,
        )
    def _submit_tasks(
        self,
        prompts: List[str],
        max_tokens: Optional[int],
        temperature: float,
        top_p: float,
        top_k: int,
        frequency_penalty: float,
        rep_window: int,
    ) -> Tuple[GenerateResult, List[str]]:
        n = len(prompts)
        result = GenerateResult(count=n)
        task_ids = []
        for i, p in enumerate(prompts):
            cb = self._make_callback(result, i)
            task_id = self.scheduler.add_task(
                prompt=p,
                max_tokens=max_tokens,
                temperature=temperature,
                top_p=top_p,
                top_k=top_k,
                frequency_penalty=frequency_penalty,
                rep_window=rep_window,
                stream_callback=cb,
            )
            task_ids.append(task_id)
        return result, task_ids
    @staticmethod
    def _make_callback(result: GenerateResult, idx: int):
        def cb(token):
            result.append(token, idx)
        return cb
    def _generate_streaming(
        self,
        prompts: List[str],
        is_batch: bool,
-        max_tokens: int,
+        max_tokens: Optional[int],
        temperature: float,
        top_p: float,
        top_k: int,
-    ) -> Generator[str, None, None]:
+        frequency_penalty: float,
-        """Internal streaming generator.
+        rep_window: int,
-
+    ) -> Generator:
-        Polls the _Result accumulator in a loop, yielding tokens as they arrive.
+        result, task_ids = self._submit_tasks(
-        Cleans up the scheduler task on GeneratorExit.
+            prompts,
-
+            max_tokens,
-        Args:
+            temperature,
-            prompts: List of prompts (only first is used; batch not yet supported).
+            top_p,
-            is_batch: If True, raises NotImplementedError.
+            top_k,
-            max_tokens: Maximum tokens to generate.
+            frequency_penalty,
-            temperature: Sampling temperature.
+            rep_window,
            top_p: Nucleus sampling threshold.
            top_k: Top-k sampling count.
        Yields:
            Decoded token strings.
        """
        if is_batch:
            raise NotImplementedError("Batch streaming not yet supported")
        result = _Result()
        task_id = self.scheduler.add_task(
            prompt=prompts[0],
            max_tokens=max_tokens,
            temperature=temperature,
            top_p=top_p,
            top_k=top_k,
            stream_callback=lambda tok: result.append(tok, 0),
        )
        n = len(prompts)
        remaining = n
        finished = [False] * n
        def gen():
            nonlocal remaining
            try:
-                while True:
+                while remaining > 0:
-                    tokens = result.pop_all()
+                    items = result.pop_all()
-                    for token in tokens:
+                    for idx, token in items:
                        if token is STOP:
-                            return
+                            if not finished[idx]:
-                        yield token
+                                finished[idx] = True
-                    if not result.wait(timeout=0.05):
+                                remaining -= 1
-                        pass
+                        else:
                            yield (idx, token) if is_batch else token
                    if remaining > 0:
                        result.wait(timeout=0.05)
            finally:
-                self.scheduler.remove_task(task_id)
+                for tid in task_ids:
                    self.scheduler.remove_task(tid)
        return gen()
@@ -398,62 +308,40 @@ class InferenceEngine:
        self,
        prompts: List[str],
        is_batch: bool,
-        max_tokens: int,
+        max_tokens: Optional[int],
        temperature: float,
        top_p: float,
        top_k: int,
        frequency_penalty: float,
        rep_window: int,
    ) -> Union[str, List[str]]:
-        """Internal non-streaming generator.
+        result, task_ids = self._submit_tasks(
-
+            prompts,
-        Submits all prompts to the scheduler and waits for all to complete.
+            max_tokens,
-
+            temperature,
-        Args:
+            top_p,
-            prompts: List of prompt strings.
+            top_k,
-            is_batch: Whether multiple prompts were provided.
+            frequency_penalty,
-            max_tokens: Maximum tokens to generate.
+            rep_window,
            temperature: Sampling temperature.
            top_p: Nucleus sampling threshold.
            top_k: Top-k sampling count.
        Returns:
            Single string for one prompt, list of strings for batch.
        """
        result = _Result(count=len(prompts))
        task_ids = []
        for i, p in enumerate(prompts):
            def make_cb(idx):
                return lambda tok: result.append(tok, idx)
            task_id = self.scheduler.add_task(
                prompt=p,
                max_tokens=max_tokens,
                temperature=temperature,
                top_p=top_p,
                top_k=top_k,
                stream_callback=make_cb(i),
        )
            task_ids.append(task_id)
        try:
            result.wait_completion()
        except TimeoutError:
            for tid in task_ids:
                self.scheduler.remove_task(tid)
            raise
-        for task_id in task_ids:
+        for tid in task_ids:
-            self.scheduler.remove_task(task_id)
+            self.scheduler.remove_task(tid)
        res = result.get_results()
        return res if is_batch else res[0]
    def get_stats(self) -> Dict[str, Any]:
        """Returns current engine statistics.
        Returns:
            Dict with total_tasks, total_tokens, active_tasks, waiting_queue.
        """
        return self.scheduler.get_stats()
-    def shutdown(self) -> None:
+    def shutdown(self):
        """Shuts down the engine, stops the scheduler, and frees GPU memory."""
        self.scheduler.stop()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
@@ -0,0 +1,343 @@
 """Composable sampling strategies for logit transformation.
 Implements the Strategy pattern: each sampling technique
 (temperature, top-k, top-p, frequency penalty) is a pluggable
 strategy that can be composed into a pipeline.
 All strategies accept both scalar and per-sample tensor
 parameters, so a single pipeline works for any batch size.
 """
 from abc import ABC, abstractmethod
 from typing import List, Optional, Union
 import torch
 from torch import Tensor
 class BaseSamplingStrategy(ABC):
    """Abstract base for a logit transformation strategy."""
    @abstractmethod
    def apply(
        self,
        logits: Tensor,
        filter_value: float = -float("inf"),
        input_ids: Optional[Tensor] = None,
        input_mask: Optional[Tensor] = None,
    ) -> Tensor:
        """Applies the strategy to logits.
        Args:
            logits: Raw logits tensor (batch, vocab_size).
            filter_value: Value assigned to filtered-out positions.
            input_ids: Previously generated token IDs ``[batch, seq_len]``,
                padded with 0. Used by frequency penalty.
            input_mask: Boolean mask ``[batch, seq_len]``, True for real
                tokens, False for padding. Used to exclude padding from
                penalty computation.
        Returns:
            Transformed logits tensor.
        """
        raise NotImplementedError
 class TemperatureStrategy(BaseSamplingStrategy):
    """Divides logits by temperature to control randomness.
    Args:
        temperature: Scalar or ``[batch]`` tensor.
    """
    def __init__(self, temperature: Union[float, Tensor] = 1.0):
        self.temperature = temperature
    def apply(
        self,
        logits: Tensor,
        filter_value: float = -float("inf"),
        input_ids: Optional[Tensor] = None,
        input_mask: Optional[Tensor] = None,
    ) -> Tensor:
        t = self.temperature
        if isinstance(t, Tensor):
            t = t.to(logits.device, non_blocking=True).view(-1, 1)
            t = torch.clamp(t, min=1e-8)
            if (t != 1.0).any():
                logits = logits / t
        elif t != 1.0:
            logits = logits / max(t, 1e-8)
        return logits
 class TopKStrategy(BaseSamplingStrategy):
    """Keeps only the top-k logits, setting the rest to filter_value.
    Args:
        top_k: Scalar or ``[batch]`` tensor (0 disables).
    """
    def __init__(self, top_k: Union[int, Tensor] = 0):
        self.top_k = top_k
    def apply(
        self,
        logits: Tensor,
        filter_value: float = -float("inf"),
        input_ids: Optional[Tensor] = None,
        input_mask: Optional[Tensor] = None,
    ) -> Tensor:
        tk = self.top_k
        if isinstance(tk, Tensor):
            tk = tk.to(logits.device, non_blocking=True).long().clamp(min=0)
            max_k = int(tk.max().item())
            if max_k <= 0:
                return logits
            max_k = min(max_k, logits.size(-1))
            values, _ = torch.topk(logits, max_k, dim=-1)
            per_row_k = tk.clamp(max=max_k)
            thresholds = torch.full_like(logits[..., -1:], -float("inf"))
            positive = per_row_k > 0
            if positive.any():
                row_idx = torch.arange(logits.size(0), device=logits.device)[positive]
                thresholds[positive] = values[
                    row_idx, per_row_k[positive] - 1
                ].unsqueeze(-1)
            logits[logits < thresholds] = filter_value
            return logits
        if tk > 0:
            k = min(tk, logits.size(-1))
            thresholds = torch.topk(logits, k, dim=-1)[0][..., -1:]
            logits[logits < thresholds] = filter_value
        return logits
 class TopPStrategy(BaseSamplingStrategy):
    """Nucleus (top-p) filtering: keeps the smallest set of tokens whose
    cumulative probability exceeds top_p.
    Args:
        top_p: Scalar or ``[batch]`` tensor (1.0 disables).
    """
    def __init__(self, top_p: Union[float, Tensor] = 1.0):
        self.top_p = top_p
    def _apply(
        self, logits: Tensor, top_p: Union[float, Tensor], filter_value: float
    ) -> Tensor:
        sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
        cum_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
        remove = cum_probs > top_p
        remove[..., 1:] = remove[..., :-1].clone()
        remove[..., 0] = False
        mask = torch.zeros_like(logits, dtype=torch.bool)
        mask.scatter_(1, sorted_indices, remove)
        logits[mask] = filter_value
        return logits
    def apply(
        self,
        logits: Tensor,
        filter_value: float = -float("inf"),
        input_ids: Optional[Tensor] = None,
        input_mask: Optional[Tensor] = None,
    ) -> Tensor:
        tp = self.top_p
        if isinstance(tp, Tensor):
            tp = tp.to(logits.device, non_blocking=True)
            if (tp < 1.0).any():
                logits = self._apply(logits, tp.view(-1, 1), filter_value)
        elif tp < 1.0:
            logits = self._apply(logits, tp, filter_value)
        return logits
 class FrequencyPenaltyStrategy(BaseSamplingStrategy):
    """Penalizes tokens based on how many times they appeared in history.
    Subtracts ``penalty * count(token)`` from each token's logit, where
    ``count(token)`` is the number of occurrences in the generation history
    (prompt + output). A penalty of ``0.0`` disables the strategy.
    Unlike repetition penalty (which only checks *presence*), frequency
    penalty scales linearly with occurrence count: the first use is
    penalized once, the third use three times. This allows natural
    repetition of common words while suppressing degenerate loops.
    Reference: OpenAI API ``frequency_penalty`` parameter.
    Args:
        penalty: Scalar or ``[batch]`` tensor (0.0 disables, range -2.0~2.0).
    """
    def __init__(self, penalty: Union[float, Tensor] = 0.0):
        self.penalty = penalty
    def apply(
        self,
        logits: Tensor,
        filter_value: float = -float("inf"),
        input_ids: Optional[Tensor] = None,
        input_mask: Optional[Tensor] = None,
    ) -> Tensor:
        if input_ids is None:
            return logits
        p = self.penalty
        if isinstance(p, Tensor):
            p = p.to(logits.device, non_blocking=True).view(-1, 1)
            if (p == 0.0).all():
                return logits
        elif p == 0.0:
            return logits
        input_ids = input_ids.to(logits.device, non_blocking=True)
        if input_mask is not None:
            input_mask = input_mask.to(logits.device, non_blocking=True)
            masked_ids = input_ids.clone()
            masked_ids[~input_mask] = -1
        else:
            masked_ids = input_ids
        batch_sz, seq_len = masked_ids.shape
        vocab_size = logits.size(-1)
        if isinstance(p, Tensor):
            penalty_per_row = p.expand(batch_sz, 1)
        else:
            penalty_per_row = torch.full(
                (batch_sz, 1), float(p), device=logits.device, dtype=logits.dtype
            )
        counts = torch.zeros(
            batch_sz, vocab_size, device=logits.device, dtype=logits.dtype
        )
        valid_mask = masked_ids >= 0
        if valid_mask.any():
            valid_ids = masked_ids[valid_mask]
            row_indices = (
                torch.arange(batch_sz, device=logits.device)
                .unsqueeze(1)
                .expand_as(masked_ids)[valid_mask]
            )
            counts.index_put_(
                (row_indices, valid_ids),
                torch.ones_like(valid_ids, dtype=logits.dtype),
                accumulate=True,
            )
        return logits - penalty_per_row * counts
 class SamplingPipeline(BaseSamplingStrategy):
    """Composes multiple sampling strategies into a single transformation.
    Strategies are applied sequentially in the order they are provided,
    matching the original temperature -> top-k -> top-p ordering.
    Usage::
        pipeline = SamplingPipeline([
            TemperatureStrategy(0.8),
            TopKStrategy(50),
            TopPStrategy(0.95),
        ])
        logits = pipeline.apply(logits)
        token = pipeline.sample(logits)       # softmax + multinomial
    """
    def __init__(self, strategies: List[BaseSamplingStrategy]):
        self.strategies = strategies
    def apply(
        self,
        logits: Tensor,
        filter_value: float = -float("inf"),
        input_ids: Optional[Tensor] = None,
        input_mask: Optional[Tensor] = None,
    ) -> Tensor:
        for strategy in self.strategies:
            logits = strategy.apply(logits, filter_value, input_ids, input_mask)
        return logits
    @staticmethod
    def _is_greedy(temperature: Union[float, Tensor]) -> bool:
        if isinstance(temperature, Tensor):
            return temperature.numel() == 1 and temperature.item() == 0
        return temperature == 0
    @torch.inference_mode()
    def sample(
        self,
        logits: Tensor,
        filter_value: float = -float("inf"),
        input_ids: Optional[Tensor] = None,
        input_mask: Optional[Tensor] = None,
    ) -> Tensor:
        """Apply strategies then sample (softmax + multinomial).
        Short-circuits to ``argmax`` when temperature is exactly 0
        (deterministic / greedy decode).
        Args:
            logits: Raw logits ``[batch, vocab_size]``.
            input_ids: Previously generated token IDs ``[batch, seq_len]``.
            input_mask: Boolean mask for ``input_ids`` padding.
        Returns:
            Sampled token IDs ``[batch]``.
        """
        for s in self.strategies:
            if isinstance(s, TemperatureStrategy) and self._is_greedy(s.temperature):
                return logits.argmax(dim=-1)
            break
        return torch.multinomial(
            torch.softmax(
                self.apply(logits, filter_value, input_ids, input_mask), dim=-1
            ),
            num_samples=1,
        ).squeeze(-1)
@torch.inference_mode()
 def sample(
    logits: Tensor,
    temperature: Union[float, Tensor] = 1.0,
    top_k: Union[int, Tensor] = 0,
    top_p: Union[float, Tensor] = 1.0,
    frequency_penalty: Union[float, Tensor] = 0.0,
    input_ids: Optional[Tensor] = None,
    input_mask: Optional[Tensor] = None,
    filter_value: float = -float("inf"),
 ) -> Tensor:
    """Apply sampling strategies then sample (softmax + multinomial).
    Shortcut for ``SamplingPipeline(...).sample(logits)``.
    When **temperature** is exactly 0 (scalar or single-element tensor)
    the function short-circuits to ``argmax`` for deterministic decode.
    Args:
        logits: Raw logits ``[batch, vocab_size]``.
        frequency_penalty: Penalty per occurrence for repeated tokens
            (0.0 disables, range -2.0~2.0).
        input_ids: Previously generated token IDs ``[batch, seq_len]``.
        input_mask: Boolean mask for ``input_ids`` padding.
    Returns:
        Sampled token IDs ``[batch]``.
    """
    if SamplingPipeline._is_greedy(temperature):
        return logits.argmax(dim=-1)
    return SamplingPipeline(
        [
            TemperatureStrategy(temperature),
            TopKStrategy(top_k),
            TopPStrategy(top_p),
            FrequencyPenaltyStrategy(frequency_penalty),
        ]
    ).sample(logits, filter_value, input_ids, input_mask)
@@ -1,178 +0,0 @@
 """Composable sampling strategies for logit transformation.
 Implements the Strategy pattern: each sampling technique
 (temperature, top-k, top-p) is a pluggable strategy that
 can be composed into a pipeline.
 All strategies accept both scalar and per-sample tensor
 parameters, so a single pipeline works for any batch size.
 """
 from abc import ABC, abstractmethod
 from typing import List, Union
 import torch
 from torch import Tensor
 class BaseSamplingStrategy(ABC):
    """Abstract base for a logit transformation strategy."""
    @abstractmethod
    def apply(self, logits: Tensor, filter_value: float = -float("inf")) -> Tensor:
        """Applies the strategy to logits.
        Args:
            logits: Raw logits tensor (batch, vocab_size).
            filter_value: Value assigned to filtered-out positions.
        Returns:
            Transformed logits tensor.
        """
 class TemperatureStrategy(BaseSamplingStrategy):
    """Divides logits by temperature to control randomness.
    Args:
        temperature: Scalar or ``[batch]`` tensor.
    """
    def __init__(self, temperature: Union[float, Tensor] = 1.0):
        self.temperature = temperature
    def apply(self, logits, filter_value=-float("inf")):
        t = self.temperature
        if isinstance(t, Tensor):
            if (t != 1.0).any():
                logits = logits / t.to(logits.device, non_blocking=True).view(-1, 1)
        elif t != 1.0:
            logits = logits / t
        return logits
 class TopKStrategy(BaseSamplingStrategy):
    """Keeps only the top-k logits, setting the rest to filter_value.
    Args:
        top_k: Scalar or ``[batch]`` tensor (0 disables).
    """
    def __init__(self, top_k: Union[int, Tensor] = 0):
        self.top_k = top_k
    def apply(self, logits, filter_value=-float("inf")):
        tk = self.top_k
        if isinstance(tk, Tensor):
            max_k = int(tk.max().item())
            if max_k <= 0:
                return logits
            k = min(max_k, logits.size(-1))
        elif tk > 0:
            k = min(tk, logits.size(-1))
        else:
            return logits
        thresholds = torch.topk(logits, k, dim=-1)[0][..., -1:]
        logits[logits < thresholds] = filter_value
        return logits
 class TopPStrategy(BaseSamplingStrategy):
    """Nucleus (top-p) filtering: keeps the smallest set of tokens whose
    cumulative probability exceeds top_p.
    Args:
        top_p: Scalar or ``[batch]`` tensor (1.0 disables).
    """
    def __init__(self, top_p: Union[float, Tensor] = 1.0):
        self.top_p = top_p
    def _apply(self, logits, top_p, filter_value):
        sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
        cum_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
        remove = cum_probs > top_p
        remove[..., 1:] = remove[..., :-1].clone()
        remove[..., 0] = False
        mask = torch.zeros_like(logits, dtype=torch.bool)
        mask.scatter_(1, sorted_indices, remove)
        logits[mask] = filter_value
        return logits
    def apply(self, logits, filter_value=-float("inf")):
        tp = self.top_p
        if isinstance(tp, Tensor):
            tp = tp.to(logits.device, non_blocking=True)
            if (tp < 1.0).any():
                logits = self._apply(logits, tp.view(-1, 1), filter_value)
        elif tp < 1.0:
            logits = self._apply(logits, tp, filter_value)
        return logits
 class SamplingPipeline(BaseSamplingStrategy):
    """Composes multiple sampling strategies into a single transformation.
    Strategies are applied sequentially in the order they are provided,
    matching the original temperature -> top-k -> top-p ordering.
    Usage::
        pipeline = SamplingPipeline([
            TemperatureStrategy(0.8),
            TopKStrategy(50),
            TopPStrategy(0.95),
        ])
        logits = pipeline.apply(logits)
        token = pipeline.sample(logits)       # softmax + multinomial
    """
    def __init__(self, strategies: List[BaseSamplingStrategy]):
        self.strategies = strategies
    def apply(self, logits, filter_value=-float("inf")):
        for strategy in self.strategies:
            logits = strategy.apply(logits, filter_value)
        return logits
    @torch.no_grad()
    def sample(self, logits: Tensor, filter_value: float = -float("inf")) -> Tensor:
        """Apply strategies then sample (softmax + multinomial).
        Args:
            logits: Raw logits ``[batch, vocab_size]``.
        Returns:
            Sampled token IDs ``[batch]``.
        """
        return torch.multinomial(
            torch.softmax(self.apply(logits, filter_value), dim=-1),
            num_samples=1,
        ).squeeze(-1)
@torch.inference_mode()
 def sample(
    logits: Tensor,
    temperature: Union[float, Tensor] = 1.0,
    top_k: Union[int, Tensor] = 0,
    top_p: Union[float, Tensor] = 1.0,
    filter_value: float = -float("inf"),
 ) -> Tensor:
    """Apply sampling strategies then sample (softmax + multinomial).
    Shortcut for ``SamplingPipeline(...).sample(logits)``.
    Args:
        logits: Raw logits ``[batch, vocab_size]``.
    Returns:
        Sampled token IDs ``[batch]``.
    """
    return SamplingPipeline(
        [
            TemperatureStrategy(temperature),
            TopKStrategy(top_k),
            TopPStrategy(top_p),
        ]
    ).sample(logits, filter_value)
@@ -1,411 +0,0 @@
 """Inference scheduler for single-GPU continuous batching with paged KV cache."""
 import logging
 import threading
 import time
 import uuid
 from enum import Enum
 from typing import Any, Callable, Dict, List, Optional, Tuple
 import torch
 from torch import Tensor
 from astrai.inference.cache import STOP, PagedCache
 from astrai.inference.sampling import sample
 from astrai.model.automodel import AutoModel
 from astrai.tokenize.tokenizer import AutoTokenizer
 logger = logging.getLogger(__name__)
 class TaskStatus(Enum):
    """Task states in the continuous batching lifecycle."""
    PENDING = "pending"
    RUNNING = "running"
    FINISHED = "finished"
    ABORTED = "aborted"
 class Task:
    """Represents a single generation request with paged KV cache tracking."""
    def __init__(
        self,
        task_id: str,
        prompt_ids: List[int],
        max_tokens: int = 1024,
        temperature: float = 1.0,
        top_p: float = 1.0,
        top_k: int = 50,
        stream_callback: Optional[Callable[[str], None]] = None,
    ):
        self.task_id = task_id
        self.prompt_ids = prompt_ids
        self.max_tokens = max_tokens
        self.temperature = temperature
        self.top_p = top_p
        self.top_k = top_k
        self.status = TaskStatus.PENDING
        self.output_ids: List[int] = []
        self.input_tokens: int = 0
        self.output_tokens: int = 0
        self.page_table: List[int] = []
        self.n_pages: int = 0
        self._prefix_cached_tokens: int = 0
        self.arrival_time = time.time()
        self.finish_time: Optional[float] = None
        self.stream_callback = stream_callback
        self._pages_freed: bool = False
    @property
    def next_pos(self) -> int:
        return self.input_tokens + len(self.output_ids)
    def is_finished(self, stop_ids: List[int]) -> bool:
        if self.output_tokens >= self.max_tokens:
            return True
        if self.output_ids and self.output_ids[-1] in stop_ids:
            return True
        return False
 class InferenceScheduler:
    """Continuous batching scheduler with paged KV cache.
    Runs a background generation loop with four phases per iteration:
      1. Cleanup finished tasks and release resources.
      2. Refill active batch from the waiting queue.
      3. Prefill newly activated tasks.
      4. Decode the largest same-position group of active tasks.
    """
    def __init__(
        self,
        model: AutoModel,
        tokenizer: AutoTokenizer,
        max_batch_size: int = 16,
        max_seq_len: Optional[int] = None,
        max_prompt_len: int = 512,
        page_size: int = 64,
        device: Optional[str] = None,
        dtype: Optional[torch.dtype] = None,
    ):
        config = model.config
        self.model = model
        self.tokenizer = tokenizer
        self.max_batch_size = max_batch_size
        self.max_seq_len = max_seq_len or config.max_len
        self.max_prompt_len = max_prompt_len
        self.page_size = page_size
        self.device = device or next(model.parameters()).device
        self.dtype = dtype or next(model.parameters()).dtype
        n_kv_heads = config.n_kv_heads
        head_dim = config.dim // config.n_heads
        n_layers = config.n_layers
        n_pages = (
            max_batch_size * (self.max_seq_len + page_size) + page_size - 1
        ) // page_size
        self.page_cache = PagedCache(
            n_layers,
            n_pages,
            page_size,
            n_kv_heads,
            head_dim,
            self.device,
            self.dtype,
        )
        self.waiting_queue: List[Task] = []
        self.active_tasks: List[Task] = []
        self._running = False
        self._task_event = threading.Event()
        self._lock = threading.Lock()
        self._total_tasks = 0
        self._total_tokens = 0
    def _n_pages_for(self, n_tokens: int) -> int:
        return (n_tokens + self.page_size - 1) // self.page_size
    def add_task(
        self,
        prompt: str,
        max_tokens: int = 1024,
        temperature: float = 1.0,
        top_p: float = 1.0,
        top_k: int = 50,
        stream_callback: Optional[Callable[[str], None]] = None,
    ) -> str:
        task_id = f"task_{int(time.time())}_{uuid.uuid4().hex[:8]}"
        prompt_ids = self.tokenizer.encode(prompt)
        if len(prompt_ids) > self.max_prompt_len:
            prompt_ids = prompt_ids[-self.max_prompt_len :]
        task = Task(
            task_id=task_id,
            prompt_ids=prompt_ids,
            max_tokens=max_tokens,
            temperature=temperature,
            top_p=top_p,
            top_k=top_k,
            stream_callback=stream_callback,
        )
        with self._lock:
            self.waiting_queue.append(task)
            self._total_tasks += 1
        self._task_event.set()
        return task_id
    def remove_task(self, task_id: str) -> None:
        with self._lock:
            removed_active = [t for t in self.active_tasks if t.task_id == task_id]
            self.waiting_queue = [t for t in self.waiting_queue if t.task_id != task_id]
            self.active_tasks = [t for t in self.active_tasks if t.task_id != task_id]
        for task in removed_active:
            if not task._pages_freed:
                self._free_pages(task.page_table)
                task.page_table.clear()
                task.n_pages = 0
                task._pages_freed = True
    def _free_pages(self, indices: List[int]) -> None:
        for idx in indices:
            self.page_cache.free(idx)
    def _record_page_hashes(self, task: Task, start_logical_page: int = 0) -> None:
        full_pages = len(task.prompt_ids) // self.page_size
        for i in range(start_logical_page, full_pages):
            self.page_cache.record_page(task.page_table[i], task.prompt_ids, i)
    def _remove_finished_tasks(self) -> None:
        finished = []
        for task in self.active_tasks:
            if task.is_finished(self.tokenizer.stop_ids):
                task.status = TaskStatus.FINISHED
                task.finish_time = time.time()
                finished.append(task)
                self._total_tokens += task.output_tokens
        for task in finished:
            if not task._pages_freed:
                self._free_pages(task.page_table)
                task.page_table.clear()
                task.n_pages = 0
                task._pages_freed = True
        self.active_tasks = [
            t for t in self.active_tasks if t.status != TaskStatus.FINISHED
        ]
    def _refill_active_batch(self) -> None:
        available = self.max_batch_size - len(self.active_tasks)
        if available <= 0:
            return
        to_add: List[Task] = []
        with self._lock:
            n = min(available, len(self.waiting_queue))
            for _ in range(n):
                to_add.append(self.waiting_queue.pop(0))
        failed: List[Task] = []
        for task in to_add:
            prompt_len = len(task.prompt_ids)
            hit_pages = self.page_cache.lookup_prefix(task.prompt_ids)
            cached_tokens = len(hit_pages) * self.page_size
            for p in hit_pages:
                self.page_cache.inc_ref(p)
            remaining = prompt_len - cached_tokens
            n_new = self._n_pages_for(remaining) if remaining > 0 else 0
            new_pages = self.page_cache.alloc_n(n_new) if n_new > 0 else []
            if remaining > 0 and not new_pages:
                for p in hit_pages:
                    self.page_cache.free(p)
                failed.append(task)
                continue
            task.page_table = hit_pages + new_pages
            task.n_pages = len(task.page_table)
            task._prefix_cached_tokens = cached_tokens
            task.status = TaskStatus.RUNNING
            self.active_tasks.append(task)
        if failed:
            with self._lock:
                self.waiting_queue[:0] = failed
    def _execute_prefill(
        self, tasks: List[Task], prompt_len: int, start_pos: int = 0
    ) -> None:
        tasks = sorted(tasks, key=lambda t: t.task_id)
        batch_sz = len(tasks)
        seq_len = prompt_len - start_pos
        input_ids = torch.empty(batch_sz, seq_len, dtype=torch.long, device=self.device)
        input_mask = torch.ones(batch_sz, seq_len, dtype=torch.bool, device=self.device)
        for i, t in enumerate(tasks):
            input_ids[i] = torch.tensor(
                t.prompt_ids[start_pos:prompt_len], device=self.device
            )
        page_tables = self._make_page_table_tensor(tasks)
        with torch.inference_mode():
            self.model(
                input_ids,
                input_mask=input_mask,
                start_pos=start_pos,
                paged_cache=self.page_cache.bind(page_tables, total_len=prompt_len),
            )
        start_logical_page = start_pos // self.page_size
        for t in tasks:
            self._record_page_hashes(t, start_logical_page=start_logical_page)
    def _execute_decode(self, tasks: List[Task], start_pos: int) -> None:
        if not tasks:
            return
        tasks = sorted(tasks, key=lambda t: t.task_id)
        batch_sz = len(tasks)
        for t in tasks:
            self._maybe_alloc_page(t, start_pos)
        input_ids = torch.tensor(
            [t.output_ids[-1] if t.output_ids else t.prompt_ids[-1] for t in tasks],
            dtype=torch.long,
            device=self.device,
        )
        active_mask = torch.ones((batch_sz, 1), dtype=torch.bool, device=self.device)
        page_tables = self._make_page_table_tensor(tasks)
        total_len = start_pos + 1
        temperatures = torch.tensor([t.temperature for t in tasks], device=self.device)
        top_ks = torch.tensor([t.top_k for t in tasks], device=self.device)
        top_ps = torch.tensor([t.top_p for t in tasks], device=self.device)
        with torch.inference_mode():
            outputs = self.model(
                input_ids.unsqueeze(1),
                input_mask=active_mask,
                paged_cache=self.page_cache.bind(page_tables, total_len=total_len),
                start_pos=start_pos,
            )
            logits = outputs["logits"][:, -1, :]
        next_tokens = sample(
            logits,
            temperature=temperatures,
            top_k=top_ks,
            top_p=top_ps,
        ).tolist()
        for t, ntok in zip(tasks, next_tokens):
            t.output_ids.append(ntok)
            t.output_tokens += 1
            pos = t.input_tokens + t.output_tokens
            self._maybe_alloc_page(t, pos)
            if t.stream_callback:
                t.stream_callback(self.tokenizer.decode([ntok]))
        for t in tasks:
            if t.is_finished(self.tokenizer.stop_ids):
                if t.stream_callback:
                    t.stream_callback(STOP)
    def _make_page_table_tensor(self, tasks: List[Task]) -> Tensor:
        max_pages = max(t.n_pages for t in tasks)
        rows = [t.page_table + [-1] * (max_pages - t.n_pages) for t in tasks]
        return torch.tensor(rows, dtype=torch.long, device=self.device)
    def _maybe_alloc_page(self, task: Task, pos: int) -> None:
        needed = self._n_pages_for(pos + 1)
        while task.n_pages < needed:
            p = self.page_cache.alloc()
            if p < 0:
                break
            task.page_table.append(p)
            task.n_pages += 1
    def _run_generation_loop(self) -> None:
        try:
            while self._running:
                self._remove_finished_tasks()
                self._refill_active_batch()
                if not self.active_tasks and not self.waiting_queue:
                    self._task_event.clear()
                    self._task_event.wait(timeout=1.0)
                    continue
                to_prefill = [t for t in self.active_tasks if t.output_tokens == 0]
                if to_prefill:
                    for t in to_prefill:
                        t.input_tokens = len(t.prompt_ids)
                    groups: Dict[Tuple[int, int], List[Task]] = {}
                    for t in to_prefill:
                        key = (len(t.prompt_ids), t._prefix_cached_tokens)
                        groups.setdefault(key, []).append(t)
                    for (prompt_len, start_pos), group in groups.items():
                        if start_pos < prompt_len:
                            self._execute_prefill(group, prompt_len, start_pos)
                pos_groups: Dict[int, List[Task]] = {}
                for t in self.active_tasks:
                    pos_groups.setdefault(t.next_pos, []).append(t)
                if pos_groups:
                    best_pos = max(pos_groups, key=lambda p: len(pos_groups[p]))
                    self._execute_decode(pos_groups[best_pos], best_pos)
        except Exception as e:
            logger.error(f"Scheduler loop crashed: {e}", exc_info=True)
            for task in self.active_tasks:
                if task.stream_callback:
                    task.stream_callback(STOP)
            for task in self.waiting_queue:
                if task.stream_callback:
                    task.stream_callback(STOP)
            raise
    def start(self) -> None:
        if not self._running:
            self._running = True
            t = threading.Thread(target=self._run_generation_loop, daemon=True)
            t.start()
            self._loop_thread = t
    def stop(self) -> None:
        self._running = False
        self._task_event.set()
        if hasattr(self, "_loop_thread"):
            self._loop_thread.join(timeout=2.0)
        self.waiting_queue.clear()
        self.active_tasks.clear()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
    def get_stats(self) -> Dict[str, Any]:
        return {
            "total_tasks": self._total_tasks,
            "total_tokens": self._total_tokens,
            "active_tasks": len(self.active_tasks),
            "waiting_queue": len(self.waiting_queue),
        }
@@ -1,486 +0,0 @@
 """
 OpenAI / Anthropic-compatible chat completion server backed by continuous-batching inference.
 """
 import json
 import logging
 import time
 import uuid
 from contextlib import asynccontextmanager
 from pathlib import Path
 from typing import Any, Dict, List, Optional, Union
 import torch
 import uvicorn
 from fastapi import FastAPI, HTTPException
 from fastapi.responses import StreamingResponse
 from pydantic import BaseModel, Field
 from astrai.inference.engine import InferenceEngine
 from astrai.model import AutoModel
 from astrai.tokenize import AutoTokenizer
 logger = logging.getLogger(__name__)
 _project_root = Path(__file__).parent.parent.parent
 class ServerState:
    def __init__(self):
        self.engine: Optional[InferenceEngine] = None
        self.config: Dict[str, Any] = {
            "device": "cuda",
            "dtype": torch.bfloat16,
            "param_path": None,
            "max_batch_size": 16,
        }
 _state = ServerState()
 class ChatMessage(BaseModel):
    role: str
    content: str
 class ChatCompletionRequest(BaseModel):
    """OpenAI Chat Completion API request body."""
    model: str = "astrai"
    messages: List[ChatMessage]
    temperature: Optional[float] = Field(default=1.0, ge=0.0, le=2.0)
    top_p: Optional[float] = Field(default=1.0, ge=0.0, le=1.0)
    top_k: Optional[int] = Field(default=50, ge=1)
    stream: Optional[bool] = False
    stop: Optional[Union[str, List[str]]] = None
    max_tokens: Optional[int] = Field(default=2048, ge=1)
    n: Optional[int] = Field(default=1, ge=1)
    presence_penalty: Optional[float] = Field(default=0.0, ge=-2.0, le=2.0)
    frequency_penalty: Optional[float] = Field(default=0.0, ge=-2.0, le=2.0)
    logit_bias: Optional[Dict[int, float]] = None
    user: Optional[str] = None
 class AnthropicMessage(BaseModel):
    role: str
    content: Union[str, List[Dict[str, Any]]]
 class MessagesRequest(BaseModel):
    """Anthropic Messages API request body."""
    model: str = "astrai"
    max_tokens: int = Field(default=1024, ge=1)
    messages: List[AnthropicMessage]
    system: Optional[str] = None
    temperature: Optional[float] = Field(default=1.0, ge=0.0, le=2.0)
    top_p: Optional[float] = Field(default=1.0, ge=0.0, le=1.0)
    top_k: Optional[int] = Field(default=50, ge=1)
    stream: Optional[bool] = False
    stop_sequences: Optional[List[str]] = None
 def configure_server(
    device: str = "cuda",
    dtype: torch.dtype = torch.bfloat16,
    param_path: Optional[Path] = None,
    max_batch_size: int = 16,
 ):
    _state.config.update(
        device=device,
        dtype=dtype,
        param_path=param_path,
        max_batch_size=max_batch_size,
    )
@asynccontextmanager
 async def lifespan(app: FastAPI):
    try:
        load_model(
            param_path=_state.config["param_path"],
            device=_state.config["device"],
            dtype=_state.config["dtype"],
            max_batch_size=_state.config["max_batch_size"],
        )
    except Exception as e:
        logger.error(f"Failed to load model: {e}")
        raise
    yield
    if _state.engine:
        _state.engine.shutdown()
        logger.info("Inference engine shutdown complete")
 app = FastAPI(title="AstrAI Inference Server", version="0.2.0", lifespan=lifespan)
 def load_model(
    param_path: Optional[Path] = None,
    device: str = "cuda",
    dtype: torch.dtype = torch.bfloat16,
    max_batch_size: int = 16,
 ):
    if param_path is None:
        param_path = _project_root / "params"
    if not param_path.exists():
        raise FileNotFoundError(f"Parameter directory not found: {param_path}")
    tokenizer = AutoTokenizer.from_pretrained(param_path)
    model = AutoModel.from_pretrained(param_path)
    model.to(device=device, dtype=dtype)
    logger.info(f"Model loaded on {device} with dtype {dtype}")
    _state.engine = InferenceEngine(
        model=model,
        tokenizer=tokenizer,
        max_batch_size=max_batch_size,
    )
    logger.info(f"Inference engine initialized with max_batch_size={max_batch_size}")
 def _get_engine() -> InferenceEngine:
    if _state.engine is None:
        raise HTTPException(status_code=503, detail="Engine not initialized")
    return _state.engine
 def _make_chunk(
    delta: Dict[str, str],
    finish_reason: Optional[str] = None,
    *,
    resp_id: str,
    created: int,
    model: str,
    index: int = 0,
 ) -> str:
    """Build a single SSE ``data:`` chunk matching OpenAI streaming format."""
    data = {
        "id": resp_id,
        "object": "chat.completion.chunk",
        "created": created,
        "model": model,
        "choices": [
            {
                "index": index,
                "delta": delta,
                "finish_reason": finish_reason,
            }
        ],
    }
    return f"data: {json.dumps(data, ensure_ascii=False)}\n\n"
@app.get("/health")
 async def health():
    return {
        "status": "ok",
        "model_loaded": _state.engine is not None,
    }
@app.get("/stats")
 async def get_stats():
    return _get_engine().get_stats()
@app.post("/v1/chat/completions")
 async def chat_completion(request: ChatCompletionRequest):
    """OpenAI-compatible chat completion endpoint (streaming + non-streaming)."""
    engine = _get_engine()
    resp_id = f"chatcmpl-{uuid.uuid4().hex[:12]}"
    created = int(time.time())
    model = request.model
    prompt = engine.tokenizer.apply_chat_template(
        [{"role": m.role, "content": m.content} for m in request.messages],
        tokenize=False,
    )
    prompt_tokens = len(engine.tokenizer.encode(prompt))
    if request.stream:
        agen = engine.generate_async(
            prompt=prompt,
            max_tokens=request.max_tokens,
            temperature=request.temperature,
            top_p=request.top_p,
            top_k=request.top_k,
        )
        async def event_stream():
            yield _make_chunk(
                {"role": "assistant"},
                finish_reason=None,
                resp_id=resp_id,
                created=created,
                model=model,
            )
            completion_tokens = 0
            async for token in agen:
                yield _make_chunk(
                    {"content": token},
                    finish_reason=None,
                    resp_id=resp_id,
                    created=created,
                    model=model,
                )
                completion_tokens += 1
            yield _make_chunk(
                {},
                finish_reason="stop",
                resp_id=resp_id,
                created=created,
                model=model,
            )
            usage = {
                "prompt_tokens": prompt_tokens,
                "completion_tokens": completion_tokens,
                "total_tokens": prompt_tokens + completion_tokens,
            }
            yield f"data: {json.dumps(usage, ensure_ascii=False)}\n\n"
            yield "data: [DONE]\n\n"
        return StreamingResponse(
            event_stream(),
            media_type="text/event-stream",
            headers={"Cache-Control": "no-cache", "Connection": "keep-alive"},
        )
    completion_tokens = 0
    chunks: List[str] = []
    agen = engine.generate_async(
        prompt=prompt,
        max_tokens=request.max_tokens,
        temperature=request.temperature,
        top_p=request.top_p,
        top_k=request.top_k,
    )
    async for token in agen:
        chunks.append(token)
        completion_tokens += 1
    content = "".join(chunks)
    return {
        "id": resp_id,
        "object": "chat.completion",
        "created": created,
        "model": model,
        "choices": [
            {
                "index": 0,
                "message": {"role": "assistant", "content": content},
                "finish_reason": "stop",
            }
        ],
        "usage": {
            "prompt_tokens": prompt_tokens,
            "completion_tokens": completion_tokens,
            "total_tokens": prompt_tokens + completion_tokens,
        },
    }
 def _make_anthropic_sse(event: str, data: Dict[str, Any]) -> str:
    return f"event: {event}\ndata: {json.dumps(data, ensure_ascii=False)}\n\n"
 def _check_stop_sequence(text: str, stop_sequences: List[str]) -> Optional[str]:
    for seq in stop_sequences:
        if seq and seq in text:
            return seq
    return None
 def _extract_text_content(content: Union[str, List[Dict[str, Any]]]) -> str:
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        for block in content:
            if isinstance(block, dict) and block.get("type") == "text":
                return block.get("text", "")
    return ""
 def _build_anthropic_messages(
    messages: List[AnthropicMessage], system: Optional[str]
 ) -> List[Dict[str, str]]:
    result: List[Dict[str, str]] = []
    if system:
        result.append({"role": "system", "content": system})
    for m in messages:
        content = _extract_text_content(m.content)
        if content:
            result.append({"role": m.role, "content": content})
    return result
@app.post("/v1/messages")
 async def create_message(request: MessagesRequest):
    """Anthropic-compatible Messages API endpoint (streaming + non-streaming)."""
    engine = _get_engine()
    resp_id = f"msg_{uuid.uuid4().hex[:24]}"
    model = request.model
    chat_messages = _build_anthropic_messages(request.messages, request.system)
    prompt = engine.tokenizer.apply_chat_template(chat_messages, tokenize=False)
    prompt_tokens = len(engine.tokenizer.encode(prompt))
    stop_sequences = request.stop_sequences or []
    if request.stream:
        agen = engine.generate_async(
            prompt=prompt,
            max_tokens=request.max_tokens,
            temperature=request.temperature,
            top_p=request.top_p,
            top_k=request.top_k,
        )
        async def event_stream():
            yield _make_anthropic_sse(
                "message_start",
                {
                    "type": "message_start",
                    "message": {
                        "id": resp_id,
                        "type": "message",
                        "role": "assistant",
                        "model": model,
                        "content": [],
                        "usage": {"input_tokens": prompt_tokens},
                    },
                },
            )
            yield _make_anthropic_sse(
                "content_block_start",
                {
                    "type": "content_block_start",
                    "index": 0,
                    "content_block": {"type": "text", "text": ""},
                },
            )
            completion_tokens = 0
            accumulated = ""
            stopped_seq: Optional[str] = None
            async for token in agen:
                accumulated += token
                completion_tokens += 1
                matched = _check_stop_sequence(accumulated, stop_sequences)
                if matched:
                    text = accumulated[: accumulated.rfind(matched)]
                    stopped_seq = matched
                    if text:
                        yield _make_anthropic_sse(
                            "content_block_delta",
                            {
                                "type": "content_block_delta",
                                "index": 0,
                                "delta": {"type": "text_delta", "text": text},
                            },
                        )
                    break
                yield _make_anthropic_sse(
                    "content_block_delta",
                    {
                        "type": "content_block_delta",
                        "index": 0,
                        "delta": {"type": "text_delta", "text": token},
                    },
                )
            yield _make_anthropic_sse(
                "content_block_stop",
                {"type": "content_block_stop", "index": 0},
            )
            stop_reason = "stop_sequence" if stopped_seq else "end_turn"
            yield _make_anthropic_sse(
                "message_delta",
                {
                    "type": "message_delta",
                    "delta": {"stop_reason": stop_reason, "stop_sequence": stopped_seq},
                    "usage": {"output_tokens": completion_tokens},
                },
            )
            yield _make_anthropic_sse(
                "message_stop",
                {"type": "message_stop"},
            )
        return StreamingResponse(
            event_stream(),
            media_type="text/event-stream",
            headers={"Cache-Control": "no-cache", "Connection": "keep-alive"},
        )
    completion_tokens = 0
    chunks: List[str] = []
    agen = engine.generate_async(
        prompt=prompt,
        max_tokens=request.max_tokens,
        temperature=request.temperature,
        top_p=request.top_p,
        top_k=request.top_k,
    )
    stopped_seq: Optional[str] = None
    accumulated = ""
    async for token in agen:
        chunks.append(token)
        completion_tokens += 1
        accumulated += token
        matched = _check_stop_sequence(accumulated, stop_sequences)
        if matched:
            stopped_seq = matched
            break
    content = "".join(chunks)
    if stopped_seq:
        idx = content.rfind(stopped_seq)
        if idx != -1:
            content = content[:idx]
    return {
        "id": resp_id,
        "type": "message",
        "role": "assistant",
        "model": model,
        "content": [{"type": "text", "text": content}],
        "stop_reason": "stop_sequence" if stopped_seq else "end_turn",
        "stop_sequence": stopped_seq,
        "usage": {
            "input_tokens": prompt_tokens,
            "output_tokens": completion_tokens,
        },
    }
 def run_server(
    host: str = "0.0.0.0",
    port: int = 8000,
    reload: bool = False,
    device: str = "cuda",
    dtype: torch.dtype = torch.bfloat16,
    param_path: Optional[Path] = None,
    max_batch_size: int = 16,
 ):
    configure_server(
        device=device,
        dtype=dtype,
        param_path=param_path,
        max_batch_size=max_batch_size,
    )
    uvicorn.run(
        "astrai.inference.server:app",
        host=host,
        port=port,
        reload=reload,
    )
@@ -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",
 ]
@@ -4,18 +4,22 @@ AutoModel base class for model loading and saving.
 from contextlib import contextmanager
 from pathlib import Path
-from typing import Self, Type, Union
+from typing import Self, Union
 import safetensors.torch as st
 import torch.nn as nn
-from astrai.config import ModelConfig
+from astrai.config.model_config import BaseModelConfig, ConfigFactory
-from astrai.factory import Registry
+from astrai.factory import BaseFactory
 from astrai.serialization import load_model_config, load_model_weights, save_model
@contextmanager
 def _disable_random_init(enable: bool = True):
-    init_functions = [
+    if not enable:
        yield
        return
    names = (
        "xavier_normal_",
        "xavier_uniform_",
        "kaiming_normal_",
@@ -25,60 +29,27 @@ def _disable_random_init(enable: bool = True):
        "constant_",
        "normal_",
        "uniform_",
-    ]
+    )
-    original_funcs = {}
+    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(nn.Module):
+class AutoModel(BaseFactory["AutoModel"], nn.Module):
    """
    Autoregressive language model base class.
-    Provides model loading/saving and generation capabilities.
+    Provides model loading/saving, registration, and generation.
    """
-    _registry = Registry()
+    def __init__(self, config: BaseModelConfig):
    def __init__(self, config: ModelConfig):
        super().__init__()
        self.config = config
    @classmethod
    def register(cls, model_type: str):
        """
        Class method decorator to register model type.
        Usage:
            @AutoModel.register('transformer')
            class Transformer(AutoModel):
                ...
        """
        def decorator(sub_cls: Type["AutoModel"]) -> Type["AutoModel"]:
            cls._registry.register(model_type.lower(), sub_cls)
            return sub_cls
        return decorator
    @classmethod
    def get_model_class(cls, model_type: str) -> Type["AutoModel"]:
        """Get model class by model_type string."""
        model_type = model_type.lower()
        if not cls._registry.contains(model_type):
            available = cls._registry.list_names()
            raise ValueError(
                f"Unknown model_type: {model_type}. Available: {available}"
            )
        return cls._registry.get(model_type)
    @classmethod
    def from_pretrained(
        cls,
@@ -89,24 +60,22 @@ class AutoModel(nn.Module):
        model_path = Path(path)
        # Load config
        config = ModelConfig()
        config_path = model_path / "config.json"
-        if config_path.exists():
+        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))
-        actual_cls = cls.get_model_class(model_type)
+        config = ConfigFactory.load(raw)
        model_type = config.model_type or "autoregressive_lm"
        actual_cls = AutoModel.get_component_class(model_type)
        with _disable_random_init(enable=disable_random_init):
            model = actual_cls(config)
        # Load weights
        weights_path = model_path / "model.safetensors"
        if weights_path.exists():
-            state_dict = st.load_file(str(weights_path))
+            state_dict = load_model_weights(str(model_path))
            model.load_state_dict(state_dict, strict=strict)
        return model
@@ -114,15 +83,12 @@ class AutoModel(nn.Module):
    def save_pretrained(
        self,
        save_directory: Union[str, Path],
-    ) -> None:
+    ):
-        save_path = Path(save_directory)
+        save_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",
 ]
@@ -0,0 +1,214 @@
 from typing import Optional
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from torch import Tensor
 from astrai.factory import BaseFactory
 from astrai.inference.core.cache import CacheView
 from astrai.model.components.linear import Linear
 from astrai.model.components.norm import RMSNorm
 from astrai.model.components.rope import apply_rotary_emb
 def repeat_kv(x: Tensor, n_rep: int) -> Tensor:
    bs, slen, n_heads, head_dim = x.shape
    if n_rep == 1:
        return x
    return (
        x[:, :, :, None, :]
        .expand(bs, slen, n_heads, n_rep, head_dim)
        .reshape(bs, slen, n_heads * n_rep, head_dim)
    )
 class AttnFactory(BaseFactory[nn.Module]):
    pass
@AttnFactory.register("gqa")
 class GQA(nn.Module):
    def __init__(
        self,
        dim: int,
        n_heads: int,
        n_kv_heads: int,
        use_qk_norm: bool,
        norm_eps: float,
        use_gated_attention: bool,
        layer_id: int,
        n_layers: int = 1,
    ):
        super().__init__()
        assert dim % n_heads == 0
        assert n_heads % n_kv_heads == 0
        self.head_dim = dim // n_heads
        self.layer_id = layer_id
        self.dim = dim
        self.n_heads = n_heads
        self.n_kv_heads = n_kv_heads
        self.n_rep = n_heads // n_kv_heads
        self.use_qk_norm = use_qk_norm
        self.use_gated_attention = use_gated_attention
        self.q_proj = Linear(dim, n_heads * self.head_dim)
        self.k_proj = Linear(dim, n_kv_heads * self.head_dim)
        self.v_proj = Linear(dim, n_kv_heads * self.head_dim)
        self.o_proj = Linear(dim, dim, init_std=0.02 / (2 * n_layers) ** 0.5)
        if self.use_qk_norm:
            self.q_norm = RMSNorm(self.head_dim, norm_eps)
            self.k_norm = RMSNorm(self.head_dim, norm_eps)
        if self.use_gated_attention:
            self.gate = Linear(dim, dim)
    def _split_heads(self, x: Tensor, n_heads) -> Tensor:
        batch_size, seq_len, _ = x.shape
        x = x.reshape(batch_size, seq_len, n_heads, self.head_dim)
        return x
    def forward(
        self,
        x: Tensor,
        rotary_emb: Tensor,
        attn_mask: Tensor = None,
        paged_cache: Optional[CacheView] = None,
    ) -> Tensor:
        is_causal = attn_mask is None
        q = self._split_heads(self.q_proj(x), self.n_heads)
        k = self._split_heads(self.k_proj(x), self.n_kv_heads)
        v = self._split_heads(self.v_proj(x), self.n_kv_heads)
        q, k = apply_rotary_emb(q, rotary_emb), apply_rotary_emb(k, rotary_emb)
        if self.use_qk_norm:
            q, k = self.q_norm(q), self.k_norm(k)
        if paged_cache is not None:
            paged_cache.write(self.layer_id, k, v)
            k, v = paged_cache.gather(self.layer_id)
        k, v = repeat_kv(k, self.n_rep), repeat_kv(v, self.n_rep)
        q, k, v = q.permute(0, 2, 1, 3), k.permute(0, 2, 1, 3), v.permute(0, 2, 1, 3)
        sdqa_out = (
            F.scaled_dot_product_attention(q, k, v, attn_mask, is_causal=is_causal)
            .permute(0, 2, 1, 3)
            .contiguous()
            .flatten(2)
        )
        if self.use_gated_attention:
            sdqa_out = sdqa_out * F.sigmoid(self.gate(x))
        out = self.o_proj(sdqa_out)
        return out
@AttnFactory.register("mla")
 class MLA(nn.Module):
    def __init__(
        self,
        dim: int,
        n_heads: int,
        n_kv_heads: int,
        kv_lora_rank: int,
        qk_nope_head_dim: int,
        qk_rope_head_dim: int,
        norm_eps: float,
        use_qk_norm: bool,
        use_gated_attention: bool,
        layer_id: int,
        n_layers: int = 1,
    ):
        super().__init__()
        self.dim = dim
        self.n_heads = n_heads
        self.n_kv_heads = n_kv_heads
        self.kv_lora_rank = kv_lora_rank
        self.qk_nope_head_dim = qk_nope_head_dim
        self.qk_rope_head_dim = qk_rope_head_dim
        self.head_dim = qk_nope_head_dim + qk_rope_head_dim
        self.layer_id = layer_id
        self.n_rep = n_heads // n_kv_heads
        self.use_qk_norm = use_qk_norm
        self.use_gated_attention = use_gated_attention
        self.q_proj = Linear(dim, n_heads * self.head_dim, bias=False)
        if self.use_qk_norm:
            self.q_norm = RMSNorm(self.head_dim, norm_eps)
            self.k_norm = RMSNorm(self.head_dim, norm_eps)
        self.kv_a_proj = Linear(dim, kv_lora_rank, bias=False)
        self.kv_norm = RMSNorm(kv_lora_rank, norm_eps)
        self.kv_b_proj = Linear(
            kv_lora_rank,
            n_kv_heads * (2 * self.head_dim),
        )
        self.o_proj = Linear(
            dim, dim, bias=False, init_std=0.02 / (2 * n_layers) ** 0.5
        )
        if use_gated_attention:
            self.gate = Linear(dim, dim, bias=False)
    def forward(
        self,
        x: Tensor,
        rotary_emb: Tensor,
        attn_mask: Tensor = None,
        paged_cache: Optional[CacheView] = None,
    ) -> Tensor:
        bsz, seq_len, _ = x.size()
        is_causal = attn_mask is None
        q = self.q_proj(x)
        q = q.view(bsz, seq_len, self.n_heads, self.head_dim)
        kv_compressed = self.kv_a_proj(x)
        kv_compressed = self.kv_norm(kv_compressed)
        kv = self.kv_b_proj(kv_compressed)
        kv = kv.view(bsz, seq_len, self.n_kv_heads, -1)
        k_nope, k_rope, v = torch.split(
            kv, [self.qk_nope_head_dim, self.qk_rope_head_dim, self.head_dim], dim=-1
        )
        q_nope, q_rope = (
            q[..., : self.qk_nope_head_dim],
            q[..., self.qk_nope_head_dim :],
        )
        q_rope = apply_rotary_emb(q_rope, rotary_emb)
        k_rope = apply_rotary_emb(k_rope, rotary_emb)
        q = torch.cat([q_nope, q_rope], dim=-1)
        k = torch.cat([k_nope, k_rope], dim=-1)
        if self.use_qk_norm:
            q = self.q_norm(q)
            k = self.k_norm(k)
        if paged_cache is not None:
            paged_cache.write(self.layer_id, k, v)
            k, v = paged_cache.gather(self.layer_id)
        q = q.permute(0, 2, 1, 3)
        k = k.permute(0, 2, 1, 3)
        v = v.permute(0, 2, 1, 3)
        attn_out = F.scaled_dot_product_attention(
            q, k, v, attn_mask, is_causal=is_causal
        )
        attn_out = attn_out.permute(0, 2, 1, 3).contiguous().flatten(2)
        if self.use_gated_attention:
            attn_out = attn_out * F.sigmoid(self.gate(x))
        out = self.o_proj(attn_out)
        return out
@@ -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,337 +0,0 @@
 from typing import Optional, Tuple
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from torch import Tensor
 from astrai.inference.cache import CacheView
 def repeat_kv(x: Tensor, n_rep: int) -> Tensor:
    """Repeat KV heads n_rep times for GQA."""
    bs, slen, n_heads, head_dim = x.shape
    if n_rep == 1:
        return x
    return (
        x[:, :, :, None, :]
        .expand(bs, slen, n_heads, n_rep, head_dim)
        .reshape(bs, slen, n_heads * n_rep, head_dim)
    )
 def get_rotary_emb(
    dim: int,
    max_len: int,
    base: float = 10000,
    device: Optional[torch.device] = None,
 ) -> Tuple[Tensor, Tensor]:
    """Precompute cos/sin for RoPE."""
    theta = base ** (-torch.arange(0, dim, 2, dtype=torch.float64, device=device) / dim)
    t = torch.arange(0, max_len, dtype=torch.float64, device=device)
    freqs = torch.outer(t, theta)
    return torch.cos(freqs).float(), torch.sin(freqs).float()
 def apply_rotary_emb(x: torch.Tensor, rotary_emb: Tuple[Tensor, Tensor]) -> Tensor:
    """Apply rotary embedding via cos/sin (shape-preserving)."""
    dtype = x.dtype
    cos, sin = rotary_emb
    cos = cos.unsqueeze(0).unsqueeze(2)
    sin = sin.unsqueeze(0).unsqueeze(2)
    x_real = x[..., 0::2]
    x_imag = x[..., 1::2]
    x_real_rot = x_real * cos - x_imag * sin
    x_imag_rot = x_real * sin + x_imag * cos
    x_out = torch.stack([x_real_rot, x_imag_rot], dim=-1)
    x_out = x_out.view(*x_out.shape[:-2], -1)
    return x_out.to(dtype)
 class RotaryEmbedding(nn.Module):
    def __init__(self, dim: int, max_len: int, base: int = 10000):
        super().__init__()
        self.dim = dim
        self.max_len = max_len
        self.base = base
        self.max_len_cached = None
        self._set_rotary_buffer(self.max_len, None)
    def _set_rotary_buffer(self, max_len: int, device: Optional[torch.device] = None):
        cos_cached, sin_cached = get_rotary_emb(self.dim, max_len, self.base, device)
        self.register_buffer("cos_cached", cos_cached, persistent=False)
        self.register_buffer("sin_cached", sin_cached, persistent=False)
        self.max_len_cached = max_len
    def forward(self, x: Tensor, start_pos: int = 0) -> Tuple[Tensor, Tensor]:
        seq_len = x.size(1)
        if self.max_len_cached < seq_len + start_pos:
            self._set_rotary_buffer(self.max_len_cached * 2, x.device)
        cos = self.cos_cached[start_pos : start_pos + seq_len]
        sin = self.sin_cached[start_pos : start_pos + seq_len]
        return (cos, sin)
 class Linear(nn.Module):
    def __init__(self, in_dim: int, out_dim: int, bias: bool = False):
        super().__init__()
        self.weight = nn.Parameter(torch.empty((out_dim, in_dim)))
        self.bias = nn.Parameter(torch.zeros(out_dim)) if bias else None
    def forward(self, x: Tensor) -> Tensor:
        return F.linear(x, self.weight, self.bias)
 class RMSNorm(nn.Module):
    def __init__(self, dim, norm_eps):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(dim))
        self.normalized_shape = (dim,)
        self.norm_eps = norm_eps
    def forward(self, x: Tensor) -> Tensor:
        return F.rms_norm(x, self.normalized_shape, self.weight, self.norm_eps)
 class MLP(nn.Module):
    def __init__(self, dim: int, dim_feed_forward: int):
        super().__init__()
        self.up = Linear(dim, dim_feed_forward)
        self.gate = Linear(dim, dim_feed_forward)
        self.down = Linear(dim_feed_forward, dim)
    def forward(self, x: Tensor) -> Tensor:
        gated = self.up(x) * F.silu(self.gate(x))
        out = self.down(gated)
        return out
 class GQA(nn.Module):
    def __init__(
        self,
        dim: int,
        n_heads: int,
        n_kv_heads: int,
        use_qk_norm: bool,
        norm_eps: float,
        use_gated_attention: bool,
        layer_id: int,
    ):
        super().__init__()
        assert dim % n_heads == 0
        assert n_heads % n_kv_heads == 0
        self.head_dim = dim // n_heads
        self.layer_id = layer_id
        self.dim = dim
        self.n_heads = n_heads
        self.n_kv_heads = n_kv_heads
        self.n_rep = n_heads // n_kv_heads
        self.use_qk_norm = use_qk_norm
        self.use_gated_attention = use_gated_attention
        self.q_proj = Linear(dim, n_heads * self.head_dim)
        self.k_proj = Linear(dim, n_kv_heads * self.head_dim)
        self.v_proj = Linear(dim, n_kv_heads * self.head_dim)
        self.o_proj = Linear(dim, dim)
        if self.use_qk_norm:
            self.q_norm = RMSNorm(self.head_dim, norm_eps)
            self.k_norm = RMSNorm(self.head_dim, norm_eps)
        if self.use_gated_attention:
            self.gate = Linear(dim, dim)
    def _split_heads(self, x: Tensor, n_heads) -> Tensor:
        batch_size, seq_len, _ = x.shape
        x = x.reshape(batch_size, seq_len, n_heads, self.head_dim)
        return x
    def forward(
        self,
        x: Tensor,
        rotary_emb: Tuple[Tensor, Tensor],
        mask: Tensor = None,
        paged_cache: Optional[CacheView] = None,
        start_pos: int = 0,
    ) -> Tensor:
        bsz, seq_len, _ = x.size()
        is_causal = mask is None
        # (bsz, seq_len, dim) -> (bsz, seq_len, n_heads, head_dim)
        q = self._split_heads(self.q_proj(x), self.n_heads)
        k = self._split_heads(self.k_proj(x), self.n_kv_heads)
        v = self._split_heads(self.v_proj(x), self.n_kv_heads)
        q, k = apply_rotary_emb(q, rotary_emb), apply_rotary_emb(k, rotary_emb)
        if self.use_qk_norm:
            q, k = self.q_norm(q), self.k_norm(k)
        if paged_cache is not None:
            paged_cache.write(self.layer_id, start_pos, k, v)
            k, v = paged_cache.gather(self.layer_id)
        k, v = repeat_kv(k, self.n_rep), repeat_kv(v, self.n_rep)
        # (bsz, seq_len, n_heads, head_dim) -> (bsz, n_heads, seq_len, head_dim)
        q, k, v = q.permute(0, 2, 1, 3), k.permute(0, 2, 1, 3), v.permute(0, 2, 1, 3)
        sdqa_out = (
            F.scaled_dot_product_attention(q, k, v, mask, is_causal=is_causal)
            .permute(0, 2, 1, 3)
            .contiguous()
            .flatten(2)
        )
        if self.use_gated_attention:
            sdqa_out = sdqa_out * F.sigmoid(self.gate(x))
        out = self.o_proj(sdqa_out)
        return out
 class MLA(nn.Module):
    def __init__(
        self,
        dim: int,
        n_heads: int,
        n_kv_heads: int,
        kv_lora_rank: int,
        qk_nope_head_dim: int,
        qk_rope_head_dim: int,
        norm_eps: float,
        use_gated_attention: bool,
        layer_id: int,
    ):
        super().__init__()
        self.dim = dim
        self.n_heads = n_heads
        self.n_kv_heads = n_kv_heads
        self.kv_lora_rank = kv_lora_rank
        self.qk_nope_head_dim = qk_nope_head_dim
        self.qk_rope_head_dim = qk_rope_head_dim
        self.head_dim = qk_nope_head_dim + qk_rope_head_dim
        self.layer_id = layer_id
        self.n_rep = n_heads // n_kv_heads
        self.use_gated_attention = use_gated_attention
        self.q_proj = Linear(dim, n_heads * self.head_dim, bias=False)
        self.kv_a_proj = Linear(dim, kv_lora_rank, bias=False)
        self.kv_norm = RMSNorm(kv_lora_rank, norm_eps)
        # fused KV: (k_nope, k_rope, v)
        self.kv_b_proj = Linear(
            kv_lora_rank,
            n_kv_heads * (self.head_dim + qk_rope_head_dim + self.head_dim),
        )
        self.o_proj = Linear(dim, dim, bias=False)
        if use_gated_attention:
            self.gate = Linear(dim, dim, bias=False)
    def forward(
        self,
        x: Tensor,
        rotary_emb: Tuple[Tensor, Tensor],
        mask: Tensor = None,
        paged_cache: Optional[CacheView] = None,
        start_pos: int = 0,
    ) -> Tensor:
        bsz, seq_len, _ = x.size()
        is_causal = mask is None
        q = self.q_proj(x)
        q = q.view(bsz, seq_len, self.n_heads, self.head_dim)
        kv_compressed = self.kv_a_proj(x)
        kv_compressed = self.kv_norm(kv_compressed)
        kv = self.kv_b_proj(kv_compressed)
        kv = kv.view(bsz, seq_len, self.n_kv_heads, -1)
        k_nope, k_rope, v = torch.split(
            kv, [self.qk_nope_head_dim, self.qk_rope_head_dim, self.head_dim], dim=-1
        )
        q_nope, q_rope = (
            q[..., : self.qk_nope_head_dim],
            q[..., self.qk_rope_head_dim :],
        )
        q_rope = apply_rotary_emb(q_rope, rotary_emb)
        k_rope = apply_rotary_emb(k_rope, rotary_emb)
        q = torch.cat([q_nope, q_rope], dim=-1)
        k = torch.cat([k_nope, k_rope], dim=-1)
        if paged_cache is not None:
            paged_cache.write(self.layer_id, start_pos, k, v)
            k, v = paged_cache.gather(self.layer_id)
        q = q.permute(0, 2, 1, 3)
        k = k.permute(0, 2, 1, 3)
        v = v.permute(0, 2, 1, 3)
        attn_out = F.scaled_dot_product_attention(q, k, v, mask, is_causal=is_causal)
        attn_out = attn_out.permute(0, 2, 1, 3).contiguous().flatten(2)
        if self.use_gated_attention:
            attn_out = attn_out * F.sigmoid(self.gate(x))
        out = self.o_proj(attn_out)
        return out
 class DecoderBlock(nn.Module):
    def __init__(
        self,
        dim: int,
        n_heads: int,
        dim_ffn: int,
        n_kv_heads: int,
        norm_eps: int,
        use_qk_norm: bool,
        use_gated_attention: bool,
        layer_id: int,
    ):
        super().__init__()
        self.attention = GQA(
            dim,
            n_heads,
            n_kv_heads,
            use_qk_norm,
            norm_eps,
            use_gated_attention,
            layer_id,
        )
        self.input_norm = RMSNorm(dim, norm_eps)
        self.mlp = MLP(dim, dim_ffn)
        self.post_attention_norm = RMSNorm(dim, norm_eps)
    def forward(
        self,
        x: Tensor,
        rotary_emb: Tuple[Tensor, Tensor],
        attention_mask: Optional[Tensor] = None,
        paged_cache: Optional[CacheView] = None,
        start_pos: int = 0,
    ) -> Tensor:
        attn_output = self.attention(
            self.input_norm(x),
            rotary_emb,
            attention_mask,
            paged_cache,
            start_pos,
        )
        x = attn_output + x
        x = self.mlp(self.post_attention_norm(x)) + x
        return x
 class Embedding(nn.Module):
    def __init__(self, vocab_size: int, embedding_dim: int):
        super().__init__()
        self.weight = nn.Parameter(torch.empty((vocab_size, embedding_dim)))
    def forward(self, x: Tensor) -> Tensor:
        return F.embedding(x, self.weight)
@@ -1,98 +1,83 @@
-from typing import Any, Mapping, Optional
+from typing import Any, Dict, Mapping, Optional
 import torch
 import torch.nn as nn
 from torch import Tensor
-from astrai.config.model_config import ModelConfig
+from astrai.config.model_config import AutoRegressiveLMConfig
-from astrai.inference.cache import CacheView
+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(
    seq_mask: Tensor,
    input_tensor: Tensor,
-    start_pos: int = 0,
+    position_ids: Optional[Tensor],
    input_mask: Optional[Tensor] = None,
    is_causal: bool = False,
-) -> Tensor:
+) -> Optional[Tensor]:
-    """Build 4D attention mask from 2D seq_mask, with optional causal masking."""
+    if position_ids is None:
    device = input_tensor.device
    dtype = input_tensor.dtype
    seq_len = input_tensor.size(1)
    if seq_mask is None:
        if start_pos != 0:
            seq_mask = torch.ones((1, seq_len), dtype=torch.bool, device=device)
        else:
        return None
    if input_mask is not None and input_mask.dim() > 2:
        return input_mask
-    if seq_mask.dim() > 2:
+    device = input_tensor.device
-        return seq_mask
+    B = input_tensor.size(0)
    T = position_ids.max().item() + 1
-    batch_size = seq_mask.size(0)
+    if input_mask is None:
-    seq_mask = seq_mask[:, : start_pos + seq_len].to(device=device, dtype=torch.bool)
+        if position_ids.min().item() == 0 and is_causal:
-    expanded_mask = seq_mask.unsqueeze(1).expand(
+            return None
-        batch_size, seq_len, start_pos + seq_len
+        attend = torch.ones(B, 1, T, dtype=torch.bool, device=device)
-    )
+    else:
        attend = input_mask[:, :T].to(device=device, dtype=torch.bool).unsqueeze(1)
    if is_causal:
-        expanded_mask = torch.tril(expanded_mask, diagonal=start_pos)
+        causal = position_ids.unsqueeze(-1) >= torch.arange(T, device=device)
        attend = attend & causal
-    attention_mask = torch.zeros_like(expanded_mask, dtype=dtype, device=device)
+    return attend.unsqueeze(1)
    attention_mask = attention_mask.masked_fill_(
        ~expanded_mask, -torch.finfo(dtype).max / 2
    ).unsqueeze(1)
    return attention_mask
-@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"
@@ -100,7 +85,7 @@ class Transformer(AutoModel):
        state_dict = dict(state_dict)
-        if self.config.tie_weight:
+        if self.config.tie_weight is True:
            # same tensor for embed and lm_head
            if embed_key in state_dict:
                state_dict[lm_head_key] = state_dict[embed_key]
@@ -116,7 +101,7 @@ class Transformer(AutoModel):
            destination=destination, prefix=prefix, keep_vars=keep_vars
        )
-        if self.config.tie_weight:
+        if self.config.tie_weight is True:
            lm_head_key = prefix + "lm_head.weight"
            if lm_head_key in state_dict:
                del state_dict[lm_head_key]
@@ -128,17 +113,16 @@ class Transformer(AutoModel):
        input_ids: Tensor,
        input_mask: Optional[Tensor] = None,
        paged_cache: Optional[CacheView] = None,
-        start_pos: int = 0,
+        position_ids: Optional[Tensor] = None,
-    ) -> Tensor:
+    ) -> Dict[str, Tensor]:
        assert input_ids.ndim == 2
        x = self.embed_tokens(input_ids)
-        rotary_emb = self.rotary_embedding(x, start_pos)
+        rotary_emb = self.rotary_embedding(x, position_ids)
-
+        attn_mask = process_attention_mask(x, position_ids, input_mask, is_causal=True)
        attn_mask = process_attention_mask(input_mask, x, start_pos, is_causal=True)
        for layer in self.layers:
-            x = layer(x, rotary_emb, attn_mask, paged_cache, start_pos)
+            x = layer(x, rotary_emb, attn_mask, paged_cache)
        hidden_states = self.norm(x)
        logits = self.lm_head(hidden_states)
@@ -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,311 @@
 """Unified training executor — parallel strategy + gradient accumulation."""
 import contextlib
 import logging
 import os
 from contextlib import contextmanager
 from typing import Optional, Tuple
 import torch
 import torch.distributed as dist
 import torch.nn as nn
 from torch.distributed.fsdp import FullStateDictConfig, StateDictType
 from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
 from torch.nn.parallel import DistributedDataParallel as DDP
 from torch.optim import Optimizer
 from torch.optim.lr_scheduler import LRScheduler
 from torch.utils.data import DataLoader
 from astrai.factory import BaseFactory
 from astrai.parallel.setup import get_rank, get_world_size
 logger = logging.getLogger(__name__)
 class GradientState:
    def __init__(self, grad_accum_steps: int = 1):
        self.num_steps = max(grad_accum_steps, 1)
        self._step: int = 0
        self._sync_gradients: bool = True
    @property
    def sync_gradients(self) -> bool:
        return self._sync_gradients
    def _do_sync(self):
        self._step += 1
        self._sync_gradients = self._step % self.num_steps == 0
 class AccumOptimizer:
    def __init__(self, optimizer: Optimizer, gradient_state: GradientState):
        self.optimizer = optimizer
        self.gradient_state = gradient_state
    def step(self, closure=None):
        if self.gradient_state.sync_gradients:
            self.optimizer.step(closure)
    def zero_grad(self):
        if self.gradient_state.sync_gradients:
            self.optimizer.zero_grad()
    @property
    def param_groups(self):
        return self.optimizer.param_groups
    def state_dict(self):
        return self.optimizer.state_dict()
    def load_state_dict(self, d):
        self.optimizer.load_state_dict(d)
 class AccumScheduler:
    def __init__(self, scheduler: LRScheduler, gradient_state: GradientState):
        self.scheduler = scheduler
        self.gradient_state = gradient_state
    def step(self):
        if self.gradient_state.sync_gradients:
            self.scheduler.step()
    def state_dict(self):
        return self.scheduler.state_dict()
    def load_state_dict(self, d):
        self.scheduler.load_state_dict(d)
    def get_last_lr(self):
        return self.scheduler.get_last_lr()
 class BaseExecutor:
    def __init__(self, grad_accum_steps: int = 1):
        self.gradient_state = GradientState(grad_accum_steps)
    def prepare(
        self,
        model: nn.Module,
        optimizer: Optional[Optimizer] = None,
        dataloader: Optional[DataLoader] = None,
        scheduler: Optional[LRScheduler] = None,
    ) -> Tuple[
        nn.Module, Optional[Optimizer], Optional[DataLoader], Optional[LRScheduler]
    ]:
        model = self._prepare_model(model)
        if optimizer is not None:
            optimizer = AccumOptimizer(optimizer, self.gradient_state)
        if scheduler is not None:
            scheduler = AccumScheduler(scheduler, self.gradient_state)
        return model, optimizer, dataloader, scheduler
    def _prepare_model(self, model: nn.Module) -> nn.Module:
        return model
    def _no_sync(self, model: nn.Module):
        return contextlib.nullcontext()
    @contextmanager
    def accumulate(self, model: nn.Module):
        self.gradient_state._do_sync()
        if not self.gradient_state.sync_gradients:
            with self._no_sync(model):
                yield
        else:
            yield
    def backward(self, loss: torch.Tensor):
        loss.backward()
    def unwrap_model(self, model: nn.Module):
        return model.state_dict()
    @contextmanager
    def checkpoint_context(self, model: nn.Module):
        if self.use_distributed:
            dist.barrier()
        state_dict = self._gather_state_dict(model)
        yield state_dict
        if self.use_distributed:
            dist.barrier()
    def _gather_state_dict(self, model: nn.Module):
        state_dict = self.unwrap_model(model)
        if self.use_distributed and get_rank() != 0:
            return None
        return state_dict
    @property
    def use_distributed(self) -> bool:
        return get_world_size() > 1
    @property
    def sync_gradients(self) -> bool:
        return self.gradient_state.sync_gradients
    @property
    def grad_accum_steps(self) -> int:
        return self.gradient_state.num_steps
    def clip_grad_norm(self, model: nn.Module, max_norm: Optional[float]) -> float:
        if max_norm is None:
            total_norm = torch.norm(
                torch.stack(
                    [p.grad.norm(2) for p in model.parameters() if p.grad is not None]
                )
            )
            return total_norm.item()
        total_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
        if isinstance(total_norm, torch.Tensor):
            return total_norm.item()
        return total_norm
 class ExecutorFactory(BaseFactory[BaseExecutor]):
    pass
@ExecutorFactory.register("none")
 class NoneExecutor(BaseExecutor):
    pass
@ExecutorFactory.register("ddp")
 class DDPExecutor(BaseExecutor):
    def __init__(
        self,
        grad_accum_steps: int = 1,
        dim: int = 0,
        broadcast_buffers: bool = True,
        init_sync: bool = True,
        process_group=None,
        bucket_cap_mb: int = 25,
        find_unused_parameters: bool = False,
        check_reduction: bool = False,
        gradient_as_bucket_view: bool = False,
        static_graph: bool = False,
        delay_all_reduce_named_params=None,
        param_to_hook_all_reduce=None,
        mixed_precision=None,
        device_mesh=None,
    ):
        super().__init__(grad_accum_steps=grad_accum_steps)
        self._ddp_kwargs = dict(
            dim=dim,
            broadcast_buffers=broadcast_buffers,
            init_sync=init_sync,
            process_group=process_group,
            bucket_cap_mb=bucket_cap_mb,
            find_unused_parameters=find_unused_parameters,
            check_reduction=check_reduction,
            gradient_as_bucket_view=gradient_as_bucket_view,
            static_graph=static_graph,
            delay_all_reduce_named_params=delay_all_reduce_named_params,
            param_to_hook_all_reduce=param_to_hook_all_reduce,
            mixed_precision=mixed_precision,
            device_mesh=device_mesh,
        )
    def _prepare_model(self, model: nn.Module) -> nn.Module:
        if not self.use_distributed:
            logger.warning("DDP backend selected but world_size=1, model not wrapped")
            return model
        local_rank = int(os.environ.get("LOCAL_RANK", get_rank()))
        model = DDP(
            model,
            device_ids=[local_rank],
            output_device=local_rank,
            **self._ddp_kwargs,
        )
        logger.info("Model wrapped with DDP (world_size=%d)", get_world_size())
        return model
    def _no_sync(self, model: nn.Module):
        if isinstance(model, DDP):
            return model.no_sync()
        return contextlib.nullcontext()
    def unwrap_model(self, model: nn.Module):
        if isinstance(model, DDP):
            return model.module.state_dict()
        return model.state_dict()
@ExecutorFactory.register("fsdp")
 class FSDPExecutor(BaseExecutor):
    def __init__(
        self,
        grad_accum_steps: int = 1,
        process_group=None,
        sharding_strategy=None,
        cpu_offload=None,
        auto_wrap_policy=None,
        backward_prefetch=None,
        mixed_precision=None,
        ignored_modules=None,
        param_init_fn=None,
        sync_module_states: bool = False,
        forward_prefetch: bool = False,
        limit_all_gathers: bool = True,
        ignored_states=None,
        device_mesh=None,
    ):
        super().__init__(grad_accum_steps=grad_accum_steps)
        self._fsdp_kwargs = {
            k: v
            for k, v in dict(
                process_group=process_group,
                sharding_strategy=sharding_strategy,
                cpu_offload=cpu_offload,
                auto_wrap_policy=auto_wrap_policy,
                backward_prefetch=backward_prefetch,
                mixed_precision=mixed_precision,
                ignored_modules=ignored_modules,
                param_init_fn=param_init_fn,
                sync_module_states=sync_module_states,
                forward_prefetch=forward_prefetch,
                limit_all_gathers=limit_all_gathers,
                use_orig_params=True,
                ignored_states=ignored_states,
                device_mesh=device_mesh,
            ).items()
            if v is not None
        }
        self._original_model: Optional[nn.Module] = None
    def _prepare_model(self, model: nn.Module) -> nn.Module:
        if not self.use_distributed:
            logger.warning("FSDP backend selected but world_size=1, model not wrapped")
            return model
        self._original_model = model
        device_id = torch.device("cuda", get_rank())
        model = FSDP(model, device_id=device_id, **self._fsdp_kwargs)
        logger.info("Model wrapped with FSDP (world_size=%d)", get_world_size())
        return model
    def _no_sync(self, model: nn.Module):
        if isinstance(model, FSDP):
            return model.no_sync()
        return contextlib.nullcontext()
    def clip_grad_norm(self, model: nn.Module, max_norm: Optional[float]) -> float:
        if max_norm is None:
            return super().clip_grad_norm(model, max_norm)
        if isinstance(model, FSDP) and self.use_distributed:
            total_norm = model.clip_grad_norm_(max_norm)
            if isinstance(total_norm, torch.Tensor):
                return total_norm.item()
            return total_norm
        return super().clip_grad_norm(model, max_norm)
    def unwrap_model(self, model: nn.Module):
        if isinstance(model, FSDP) and self.use_distributed:
            with FSDP.state_dict_type(
                model,
                StateDictType.FULL_STATE_DICT,
                FullStateDictConfig(offload_to_cpu=True, rank0_only=True),
            ):
                return model.state_dict()
        return model.state_dict()
@@ -1,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,10 +115,10 @@ 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,
@@ -111,51 +126,118 @@ def wrapper_spawn_func(
    ):
        func(**kwargs)
-    except Exception as e:
+
-        print(f"Error in rank {rank}: {e}")
+class LaunchStrategy(ABC):
    """Strategy for launching a function in a distributed context."""
    def __init__(
        self,
        world_size: int,
        backend: str,
        master_addr: str,
        master_port: str,
        device_type: str,
        start_method: str,
    ):
        self.world_size = world_size
        self.backend = backend
        self.master_addr = master_addr
        self.master_port = master_port
        self.device_type = device_type
        self.start_method = start_method
    @abstractmethod
    def launch(self, func: Callable, **kwargs):
        raise NotImplementedError
 class TorchrunStrategy(LaunchStrategy):
    """External orchestrator (torchrun, SLURM, K8s) — env vars pre-set."""
    def launch(self, func: Callable, **kwargs):
        rank = int(os.environ["RANK"])
        world_size = int(os.environ["WORLD_SIZE"])
        local_rank = int(os.environ.get("LOCAL_RANK", rank))
        with setup_parallel(
            rank=rank,
            world_size=world_size,
            local_rank=local_rank,
            backend=self.backend,
            master_addr=os.environ.get("MASTER_ADDR", self.master_addr),
            master_port=os.environ.get("MASTER_PORT", self.master_port),
            device_type=self.device_type,
        ):
            func(**kwargs)
 class LocalStrategy(LaunchStrategy):
    """Local launcher — single-process or mp.start_processes."""
    def launch(self, func: Callable, **kwargs):
        args = (
            self.world_size,
            self.backend,
            self.master_addr,
            self.master_port,
            self.device_type,
            func,
            kwargs,
        )
        if self.world_size == 1:
            _run_single_rank(0, *args)
            return
        ctx = mp.start_processes(
            _run_single_rank,
            args=args,
            nprocs=self.world_size,
            start_method=self.start_method,
            join=False,
        )
        try:
            while not ctx.join():
                pass
        except BaseException:
            for p in ctx.processes:
                p.terminate()
            ctx.join()
            raise
 def _detect_launcher() -> str:
    """Detect the distributed launcher from environment.
    Returns one of: "torchelastic", "torchrun", "external", "local".
    """
    if dist.is_torchelastic_launched():
        return "torchelastic"
    if "LOCAL_WORLD_SIZE" in os.environ:
        return "torchrun"
    if "RANK" in os.environ and "WORLD_SIZE" in os.environ:
        return "external"
    return "local"
 def spawn_parallel_fn(
    func: Callable,
    world_size: int,
    backend: str = "nccl",
    master_addr: str = "localhost",
-    master_port: 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",
    ]:
        if key in os.environ:
            del os.environ[key]
    if world_size == 1:
        device_id = torch.device(device_type, 0)
        os.environ["LOCAL_RANK"] = "0"
        os.environ["WORLD_SIZE"] = "1"
        os.environ["LOCAL_DEVICE"] = str(device_id)
        func(**kwargs)
        return
    wrapper_spawn_func_args = (
        world_size,
        backend,
        master_addr,
        master_port,
        device_type,
        func,
        kwargs,
        )
-
+    else:
-    mp.spawn(
+        strategy = LocalStrategy(
-        wrapper_spawn_func, nprocs=world_size, args=wrapper_spawn_func_args, join=True
+            world_size, backend, master_addr, master_port, device_type, start_method
        )
    strategy.launch(func, **kwargs)
@@ -0,0 +1,40 @@
 from astrai.preprocessing.builder import (
    BaseMaskBuilder,
    MaskBuilderFactory,
    MultiOutputMaskBuilder,
    SectionedMaskBuilder,
    SingleOutputMaskBuilder,
 )
 from astrai.preprocessing.packing import (
    PackingStrategy,
    PackingStrategyFactory,
    plan_bfd,
 )
 from astrai.preprocessing.pipeline import Pipeline, filter_by_length
 from astrai.preprocessing.position_id import (
    PositionIdStrategy,
    PositionIdStrategyFactory,
 )
 from astrai.preprocessing.transform import TokenizeTransform
 from astrai.preprocessing.writer import (
    StoreWriter,
    StoreWriterFactory,
 )
 __all__ = [
    "BaseMaskBuilder",
    "MaskBuilderFactory",
    "MultiOutputMaskBuilder",
    "PackingStrategy",
    "PackingStrategyFactory",
    "Pipeline",
    "PositionIdStrategy",
    "PositionIdStrategyFactory",
    "SectionedMaskBuilder",
    "SingleOutputMaskBuilder",
    "StoreWriter",
    "StoreWriterFactory",
    "TokenizeTransform",
    "filter_by_length",
    "plan_bfd",
 ]
@@ -0,0 +1,337 @@
 """Mask building for preprocessing pipeline.
 :class:`SectionRenderer` converts section specs into token ids and loss
 masks (template / text / value extraction).  :class:`SingleOutputMaskBuilder`
 handles single-output (SFT / pretrain), :class:`MultiOutputMaskBuilder`
 handles multi-output (DPO / GRPO), and :class:`SectionedMaskBuilder`
 orchestrates both modes as a façade.
 """
 from abc import ABC, abstractmethod
 from typing import Optional
 from astrai.factory import BaseFactory
 def _extract_domain(item: dict, domain_key: Optional[str]) -> str:
    if not domain_key:
        return "__default__"
    val = item.get(domain_key, "__default__")
    return val if isinstance(val, str) else "__default__"
 def _resolve_action(action: str, role: str, config) -> str:
    if action == "$role":
        return config.mask.get(role, config.mask_default)
    return action
 class SectionRenderer:
    """Render section specs into ``(ids, loss_mask)`` tuples."""
    def process_sections(
        self,
        item: dict,
        sections: list,
        config,
        tokenizer,
        *,
        is_top_level: bool = False,
    ):
        all_ids: list[int] = []
        loss_mask: list[int] = []
        has_template = any(s.get("template") for s in sections)
        is_text_config = not has_template and all(
            s["action"] == "train" for s in sections
        )
        if is_top_level and has_template and tokenizer.bos_token_id is not None:
            all_ids.append(tokenizer.bos_token_id)
            loss_mask.append(0)
        first_section = True
        for sec in sections:
            field = sec["field"]
            action = sec["action"]
            use_template = sec.get("template", False)
            add_special = sec.get(
                "add_special_tokens", not use_template and first_section
            )
            if use_template:
                success = self._append_template(
                    item, field, action, tokenizer, config, all_ids, loss_mask
                )
                if not success:
                    continue
            else:
                success = self._append_text(
                    item,
                    field,
                    action,
                    tokenizer,
                    add_special,
                    is_text_config,
                    config,
                    all_ids,
                    loss_mask,
                )
                if not success:
                    continue
            first_section = False
        max_len = config.preprocessing.max_seq_len
        all_ids = all_ids[:max_len]
        loss_mask = loss_mask[: len(all_ids)]
        if not all_ids:
            return None, None
        if is_top_level and has_template and len(all_ids) <= 1:
            return None, None
        return all_ids, loss_mask
    def process_list_field(self, item: dict, sections: list, config, tokenizer):
        """Tokenize a list-valued field, preserving per-element boundaries.
        Returns ``(list_of_id_lists, list_of_mask_lists)`` where each
        inner list corresponds to one element of the source list.  This
        is critical for GRPO where each response must stay a separate
        sequence so the strategy can form a ``[G, R]`` tensor.
        """
        per_item_ids: list[list[int]] = []
        per_item_masks: list[list[int]] = []
        for sec in sections:
            field = sec["field"]
            action = sec["action"]
            use_template = sec.get("template", False)
            values = item.get(field)
            if not isinstance(values, list):
                continue
            for val in values:
                ids: list[int] = []
                mask: list[int] = []
                if use_template:
                    if isinstance(val, list):
                        wrapper = {field: val}
                        self._append_template(
                            wrapper, field, action, tokenizer, config, ids, mask
                        )
                else:
                    wrapper = {field: str(val)}
                    self._append_text(
                        wrapper,
                        field,
                        action,
                        tokenizer,
                        False,
                        False,
                        config,
                        ids,
                        mask,
                    )
                if ids:
                    max_len = config.preprocessing.max_seq_len
                    ids = ids[:max_len]
                    mask = mask[: len(ids)]
                    per_item_ids.append(ids)
                    per_item_masks.append(mask)
        if not per_item_ids:
            return None, None
        return per_item_ids, per_item_masks
    @staticmethod
    def is_value_section(sections: list) -> bool:
        return len(sections) == 1 and sections[0].get("action") == "value"
    @staticmethod
    def extract_raw_value(item: dict, sections: list):
        sec = sections[0]
        field = sec["field"]
        raw = item.get(field)
        if raw is None:
            return None
        if isinstance(raw, list):
            return [float(v) for v in raw]
        return [float(raw)]
    def _append_template(
        self, item, field, action, tokenizer, config, all_ids, loss_mask
    ):
        messages = item.get(field)
        if not isinstance(messages, list) or not messages:
            return False
        for msg in messages:
            role = msg.get("role", "")
            act = _resolve_action(action, role, config)
            rendered = tokenizer.apply_chat_template(
                [msg], tokenize=False, add_generation_prompt=False
            )
            ids = tokenizer.encode(rendered, add_special_tokens=False)
            all_ids.extend(ids)
            val = 1 if act == "train" else 0
            loss_mask.extend([val] * len(ids))
        return True
    def _append_text(
        self,
        item,
        field,
        action,
        tokenizer,
        add_special,
        is_text_config,
        config,
        all_ids,
        loss_mask,
    ):
        text = str(item.get(field, ""))
        if not text.strip():
            return False
        if is_text_config:
            pp = config.preprocessing
            if pp.min_chars > 0 and len(text) < pp.min_chars:
                return False
            if len(text) > pp.max_chars:
                return False
        ids = tokenizer.encode(text, add_special_tokens=add_special)
        all_ids.extend(ids)
        val = 1 if action == "train" else 0
        loss_mask.extend([val] * len(ids))
        return True
 class BaseMaskBuilder(ABC):
    """Convert a JSONL item into token ids and optional loss_mask."""
    @abstractmethod
    def build(self, item: dict, config, tokenizer) -> Optional[dict]: ...
 class MaskBuilderFactory(BaseFactory["BaseMaskBuilder"]):
    pass
@MaskBuilderFactory.register("single")
 class SingleOutputMaskBuilder(BaseMaskBuilder):
    """Build a single output sequence with optional loss mask.
    Expects ``config.input.sections`` (list of section specs).
    """
    def __init__(self, renderer: Optional[SectionRenderer] = None):
        self.renderer = renderer or SectionRenderer()
    def build(self, item: dict, config, tokenizer) -> Optional[dict]:
        sections = config.input.sections
        if not sections:
            return None
        ids, mask = self.renderer.process_sections(
            item, sections, config, tokenizer, is_top_level=True
        )
        if ids is None:
            return None
        result: dict = {
            "sequence": ids,
            "domain": _extract_domain(item, config.output.domain_key),
        }
        if not all(m == 1 for m in mask):
            result["loss_mask"] = mask
        return result
@MaskBuilderFactory.register("multi")
 class MultiOutputMaskBuilder(BaseMaskBuilder):
    """Build multiple output sequences (DPO / GRPO).
    Expects ``config.input.sources`` (dict of output_key → spec).
    """
    def __init__(self, renderer: Optional[SectionRenderer] = None):
        self.renderer = renderer or SectionRenderer()
    def build(self, item: dict, config, tokenizer) -> Optional[dict]:
        sources_spec = getattr(config.input, "sources", None)
        if not sources_spec:
            return None
        result: dict = {}
        any_output = False
        for output_key, spec in sources_spec.items():
            sections = spec.get("sections", [])
            if not sections:
                continue
            if self.renderer.is_value_section(sections):
                ids = self.renderer.extract_raw_value(item, sections)
                if ids is None:
                    continue
                result[output_key] = ids
                any_output = True
                continue
            list_field = spec.get("list_field", False)
            mask_key = spec.get("mask_key", f"{output_key}_mask")
            if list_field:
                ids, mask = self.renderer.process_list_field(
                    item, sections, config, tokenizer
                )
                if ids is None:
                    continue
                # ids is List[List[int]] — preserve per-response structure
                result[output_key] = ids
                if mask is not None:
                    result[mask_key] = mask
                any_output = True
                continue
            ids, mask = self.renderer.process_sections(
                item, sections, config, tokenizer, is_top_level=True
            )
            if ids is None:
                continue
            result[output_key] = ids
            if not all(m == 1 for m in mask):
                result[mask_key] = mask
            elif "mask_key" in spec:
                result[mask_key] = mask
            any_output = True
        if not any_output:
            return None
        result["domain"] = _extract_domain(item, config.output.domain_key)
        return result
@MaskBuilderFactory.register("sectioned")
 class SectionedMaskBuilder(BaseMaskBuilder):
    """Façade that dispatches to SingleOutputMaskBuilder or MultiOutputMaskBuilder.
    Preserves backward compatibility for existing configs and code that rely
    on the ``"sectioned"`` factory name.
    """
    def __init__(self):
        self._single = SingleOutputMaskBuilder()
        self._multi = MultiOutputMaskBuilder()
    def build(self, item: dict, config, tokenizer) -> Optional[dict]:
        sources_spec = getattr(config.input, "sources", None)
        if sources_spec:
            return self._multi.build(item, config, tokenizer)
        return self._single.build(item, config, tokenizer)
@@ -0,0 +1,124 @@
 """Shared preprocessing kernel used by both :class:`Pipeline` and
 :class:`TokenizeTransform`.
 The two entry points previously duplicated ~60 % of their logic:
 record iteration, mask-builder invocation, primary-id extraction,
 per-key accumulation, dtype inference and position-id generation.
 This module factors out the common core as pure functions so that
 the online (``TokenizeTransform``) and offline (``Pipeline``) paths
 stay in lockstep.
 """
 from itertools import chain
 from typing import Dict, Iterator, List, Optional
 import torch
 from astrai.config.preprocess_config import PipelineConfig
 from astrai.preprocessing.builder import MaskBuilderFactory
 from astrai.preprocessing.position_id import PositionIdStrategyFactory
 from astrai.tokenize import AutoTokenizer
 def build_preprocessing_components(config: PipelineConfig, tokenizer_path: str):
    """Load tokenizer, mask builder and position-id strategy together.
    Both ``Pipeline`` and ``TokenizeTransform`` need the same triple;
    centralising the construction avoids drift (e.g. one path forgetting
    to create the position-id strategy).
    """
    tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
    mask_builder = MaskBuilderFactory.create("sectioned")
    position_strategy = PositionIdStrategyFactory.create(
        config.output.position_ids_mode
    )
    return tokenizer, mask_builder, position_strategy
 def primary_ids(result: dict) -> List[int]:
    """Return the first flat int-list value in *result*.
    Used for token counting and position-id generation when the
    primary key name is not known (DPO uses ``chosen``, GRPO uses
    ``prompts``, SFT uses ``sequence``).
    """
    for val in result.values():
        if isinstance(val, list) and val and isinstance(val[0], int):
            return val
    return []
 def infer_dtype(ids: List) -> torch.dtype:
    """Float values become float32, everything else int32."""
    if ids and isinstance(ids[0], float):
        return torch.float32
    return torch.int32
 def iter_raw_records(
    records: List[dict],
    mask_builder,
    config: PipelineConfig,
    tokenizer,
 ) -> Iterator[dict]:
    """Yield mask-builder output dicts for each record, skipping failures.
    Drops ``domain`` from the result (callers that need it should read
    it before calling this).  Each yielded dict maps a key
    (``sequence``, ``chosen``, ``responses``…) to either a flat
    ``List[int]`` or a nested ``List[List[int]]`` (GRPO responses/masks).
    """
    for item in records:
        result = mask_builder.build(item, config, tokenizer)
        if result is None:
            continue
        result.pop("domain", None)
        if not primary_ids(result):
            continue
        yield result
 def to_per_record_tensors(
    raw: Dict[str, list],
 ) -> Dict[str, List[torch.Tensor]]:
    """Convert an accumulated ``{key: [per-record ids]}`` dict to tensors.
    Handles three shapes transparently:
    - ``List[int]`` per record (``sequence``, ``chosen``…) → one tensor per record.
    - ``List[List[int]]`` per record (GRPO ``responses``/``masks``) → one
      ``List[Tensor]`` per record (nested), preserving the per-response
      boundary so downstream code can index responses individually.
    - ``List[int]`` for the whole shard (pre-packed keys) → single tensor.
    The detection mirrors the previous inline logic in
    ``Pipeline._flush`` and ``TokenizeTransform.apply``.
    """
    tensors: Dict[str, List[torch.Tensor]] = {}
    for key, ids_list in raw.items():
        if ids_list and isinstance(ids_list[0], list):
            tensors[key] = [
                [torch.tensor(sub, dtype=infer_dtype(sub)) for sub in ids]
                if ids and isinstance(ids[0], list)
                else torch.tensor(ids, dtype=infer_dtype(ids))
                for ids in ids_list
            ]
        else:
            tensors[key] = [
                torch.tensor(list(chain.from_iterable(ids_list)), dtype=torch.int32)
            ]
    return tensors
 def build_position_ids(
    sequences: List[List[int]],
    strategy,
 ) -> Optional[List[int]]:
    """Generate position ids for *sequences* using *strategy*.
    Returns ``None`` when the strategy produces no ids (e.g. ``none``
    mode), so callers can skip attaching the key instead of storing
    an empty list.
    """
    pos_ids = strategy.generate(sequences)
    return pos_ids or None
@@ -0,0 +1,176 @@
 """Sequence packing strategies for shard-level reordering and truncation.
 Each strategy receives the accumulated ``{key: [list of token lists]}``
 dict for a shard and returns a reordered / truncated version.  The
 pipeline later flattens the result into contiguous tensors.
 """
 from abc import ABC, abstractmethod
 from typing import Dict, List
 from astrai.factory import BaseFactory
 def _truncate(seq: List[int], max_len: int, mode: str) -> List[int]:
    if len(seq) <= max_len:
        return seq
    if mode == "keep_end":
        return seq[-max_len:]
    return seq[:max_len]
 def plan_bfd(
    sequences: List[List[int]], max_packed_len: int, truncation_mode: str = "keep_start"
 ) -> List[List[int]]:
    """Best-Fit Decreasing bin packing of *sequences* into bins.
    Returns a list of bins, each bin a list of original indices into
    *sequences*.  Bin capacities are respected on the *truncated*
    length of each sequence (so a sequence longer than
    *max_packed_len* counts at *max_packed_len*).
    Pure index-based so callers can apply the same plan to any
    aligned key (``loss_mask``, ``position_ids``…).
    """
    n = len(sequences)
    order = sorted(range(n), key=lambda i: len(sequences[i]), reverse=True)
    bins: List[List[int]] = []
    bin_lengths: List[int] = []
    for orig_idx in order:
        seq_len = len(_truncate(sequences[orig_idx], max_packed_len, truncation_mode))
        best_bin = None
        best_remain = max_packed_len + 1
        for i, bl in enumerate(bin_lengths):
            remain = max_packed_len - bl
            if seq_len <= remain < best_remain:
                best_remain = remain
                best_bin = i
        if best_bin is not None:
            bins[best_bin].append(orig_idx)
            bin_lengths[best_bin] += seq_len
        else:
            bins.append([orig_idx])
            bin_lengths.append(seq_len)
    return bins
 class PackingStrategy(ABC):
    """Reorder and truncate sequences within a shard."""
    @abstractmethod
    def apply(
        self,
        keys: Dict[str, List[List[int]]],
        max_packed_len: int,
        truncation_mode: str,
    ) -> Dict[str, List[List[int]]]:
        raise NotImplementedError
 class PackingStrategyFactory(BaseFactory["PackingStrategy"]):
    pass
@PackingStrategyFactory.register("simple")
 class SimplePacking(PackingStrategy):
    def apply(
        self,
        keys: Dict[str, List[List[int]]],
        max_packed_len: int,
        truncation_mode: str,
    ) -> Dict[str, List[List[int]]]:
        return {
            k: [_truncate(v, max_packed_len, truncation_mode) for v in vals]
            for k, vals in keys.items()
        }
@PackingStrategyFactory.register("bfd")
 class BFDPacking(PackingStrategy):
    """Best-Fit Decreasing bin packing.
    Assigns sequences to bins using a best-fit heuristic (sorted by
    decreasing length) and concatenates sequences within each bin into
    a single packed sequence.  Packed sequences are truncated to
    *max_packed_len* so that each packed bin fits within one context
    window during training.
    """
    def apply(
        self,
        keys: Dict[str, List[List[int]]],
        max_packed_len: int,
        truncation_mode: str,
    ) -> Dict[str, List[List[int]]]:
        sequences = keys.get("sequence", [])
        if not sequences:
            return keys
        bins = plan_bfd(sequences, max_packed_len, truncation_mode)
        packed: Dict[str, List[List[int]]] = {}
        for k, vals in keys.items():
            packed[k] = [
                _truncate(
                    self._concat_bin(vals, bin_indices),
                    max_packed_len,
                    truncation_mode,
                )
                for bin_indices in bins
            ]
        return packed
    @staticmethod
    def _concat_bin(vals: List[List[int]], indices: List[int]) -> List[int]:
        result: List[int] = []
        for i in indices:
            result.extend(vals[i])
        return result
@PackingStrategyFactory.register("bfd_split")
 class BFDSplitPacking(BFDPacking):
    """BFD packing with over-length sequences split into chunks.
    Sequences longer than *max_packed_len* are split into consecutive
    chunks of at most *max_packed_len* tokens instead of being
    truncated.  Each chunk becomes an independent sequence that enters
    BFD planning.  All keys (``loss_mask``, ``position_ids``, …) are
    split in lockstep so per-token alignment is preserved.
    Note: because each chunk is treated as a separate document, the
    second chunk of a split sequence loses the preceding context.
    """
    def apply(
        self,
        keys: Dict[str, List[List[int]]],
        max_packed_len: int,
        truncation_mode: str,
    ) -> Dict[str, List[List[int]]]:
        sequences = keys.get("sequence", [])
        if not sequences:
            return keys
        if max_packed_len <= 0:
            return super().apply(keys, max_packed_len, truncation_mode)
        split_keys = self._split_all(keys, max_packed_len)
        return super().apply(split_keys, max_packed_len, truncation_mode)
    @staticmethod
    def _split_all(
        keys: Dict[str, List[List[int]]], max_packed_len: int
    ) -> Dict[str, List[List[int]]]:
        """Split every sequence exceeding *max_packed_len* into chunks,
        applying the same chunk boundaries to all keys."""
        sequences = keys["sequence"]
        chunk_bounds = [list(range(0, len(s), max_packed_len)) for s in sequences]
        result: Dict[str, List[List[int]]] = {}
        for key, vals in keys.items():
            split_vals: List[List[int]] = []
            for val, starts in zip(vals, chunk_bounds):
                for start in starts:
                    split_vals.append(val[start : start + max_packed_len])
            result[key] = split_vals
        return result
@@ -0,0 +1,248 @@
 """Config-driven JSONL preprocessing pipeline.
 Composes a :class:`BaseMaskBuilder` (selected by ``input.type``) with
 sharding and flush to ``.h5`` / ``.bin`` storage.  Packing, position-id
 generation and storage writing are each delegated to pluggable strategies,
 dispatched by configuration keys.
 Record iteration, mask building, primary-id extraction and per-key
 accumulation are shared with :class:`TokenizeTransform` via the
 :mod:`astrai.preprocessing.core` helpers.
 """
 import json
 import logging
 import os
 from collections import defaultdict
 from itertools import chain
 from typing import Dict, List, Optional
 import torch
 import tqdm
 from astrai.config.preprocess_config import PipelineConfig
 from astrai.preprocessing.core import (
    build_preprocessing_components,
    iter_raw_records,
    primary_ids,
 )
 from astrai.preprocessing.packing import PackingStrategyFactory
 from astrai.preprocessing.writer import StoreWriterFactory
 logger = logging.getLogger(__name__)
 _STR_TO_DTYPE: dict[str, torch.dtype] = {
    "bool": torch.bool,
    "uint8": torch.uint8,
    "int8": torch.int8,
    "int16": torch.int16,
    "int32": torch.int32,
    "int64": torch.int64,
    "float16": torch.float16,
    "float32": torch.float32,
    "float64": torch.float64,
 }
 def filter_by_length(text: str, min_len: int = 50, max_len: int = 2_000_000) -> bool:
    return min_len <= len(text) <= max_len
 class Pipeline:
    """Tokenization pipeline driven by a declarative :class:`PipelineConfig`.
    Usage::
        config = PipelineConfig.from_file("sft_pipeline.json")
        Pipeline(config, ["data.jsonl"], output_dir="out", tokenizer_path="params").run()
    """
    def __init__(
        self,
        config: PipelineConfig,
        input_paths: list[str],
        output_dir: str,
        tokenizer_path: str,
    ):
        os.makedirs(output_dir, exist_ok=True)
        self.config = config
        self.paths = input_paths
        self.output_dir = output_dir
        self.tokenizer_path = tokenizer_path
        self.tokenizer, self.mask_builder, self._position_id = (
            build_preprocessing_components(config, tokenizer_path)
        )
        self._packer = PackingStrategyFactory.create(
            config.preprocessing.packing_strategy
        )
        self._writer = StoreWriterFactory.create(config.output.storage_format)
    def transform(self, item: dict) -> Optional[dict]:
        return self.mask_builder.build(item, self.config, self.tokenizer)
    def run(self):
        domains: dict = defaultdict(lambda: defaultdict(list))
        total_tokens = 0
        shard_idx: dict[str, int] = defaultdict(int)
        count = 0
        pp = self.config.preprocessing
        for item in tqdm.tqdm(
            self._iter_items(), desc="Tokenizing", unit="docs", mininterval=0.5
        ):
            if pp.max_items and count >= pp.max_items:
                break
            try:
                result = self.transform(item)
            except Exception:
                logger.warning(
                    "Failed to process item #%d, skipping", count + 1, exc_info=True
                )
                continue
            if result is None:
                continue
            domain = result.pop("domain", "__default__")
            ids = primary_ids(result)
            if not ids:
                continue
            bucket = domains[domain]
            self._align_bucket(bucket, result, ids)
            for key, val in result.items():
                bucket[key].append(val)
            count += 1
            total_tokens += len(ids)
            if total_tokens >= self.config.output.max_tokens_per_shard:
                self._flush(domains, shard_idx)
                domains.clear()
                total_tokens = 0
        if total_tokens > 0:
            self._flush(domains, shard_idx)
    @staticmethod
    def _align_bucket(bucket: dict, result: dict, ids: list):
        """Pad previously-accumulated keys that are missing from *result*."""
        for key in list(bucket.keys()):
            if key in result:
                continue
            bucket[key].append([0] * len(ids))
    def _iter_items(self):
        for path in self.paths:
            with open(path, "r", encoding="utf-8") as f:
                if path.endswith(".json"):
                    data = json.load(f)
                    if isinstance(data, dict):
                        yield data
                    elif isinstance(data, list):
                        yield from data
                else:
                    for line in f:
                        line = line.strip()
                        if not line:
                            continue
                        yield json.loads(line)
    def _flush(self, domains, shard_idx):
        for domain, keys in domains.items():
            idx = shard_idx[domain]
            pp = self.config.preprocessing
            original_sequences = keys.get("sequence", [])
            mode = self.config.output.position_ids_mode
            keys = self._inject_doc_reset_position_ids(keys, mode, original_sequences)
            keys = self._packer.apply(dict(keys), pp.max_packed_len, pp.truncation_mode)
            tensors = self._to_tensors(keys)
            tensors = self._inject_continuous_position_ids(
                tensors, mode, keys.get("sequence", [])
            )
            self._writer.save(self.output_dir, domain, idx, tensors)
            shard_idx[domain] = idx + 1
            first_key = "sequence" if "sequence" in tensors else next(iter(tensors))
            tqdm.tqdm.write(
                f"  saved {domain}/shard_{idx:04d}  "
                f"({tensors[first_key][0].numel():,} tokens)"
            )
    def _inject_doc_reset_position_ids(
        self,
        keys: Dict[str, list],
        mode: str,
        original_sequences: List[List[int]],
    ) -> Dict[str, list]:
        """Attach per-document position_ids before packing (``doc_reset``).
        ``doc_reset`` position ids must enter the packer so that each
        packed bin concatenates the per-doc ranges in bin order.  The
        per-record structure ``[range(len(s)) for s in seqs]`` is required
        by the packer (it concatenates per-record lists per bin); the
        ``PositionIdStrategy.generate`` flattens, so it cannot be used
        directly here — it is only consulted for the ``continuous``
        post-packing path.
        """
        if mode != "doc_reset" or not original_sequences:
            return keys
        keys["position_ids"] = [list(range(len(s))) for s in original_sequences]
        return keys
    def _inject_continuous_position_ids(
        self,
        tensors: Dict[str, List[torch.Tensor]],
        mode: str,
        packed_sequences: List[List[int]],
    ) -> Dict[str, List[torch.Tensor]]:
        """Attach a single continuous position_ids tensor after packing.
        ``continuous`` mode spans the whole shard (post-packing), so it
        cannot participate in bin packing — it is computed from the
        packed sequences and appended directly to the tensor dict.
        """
        if mode != "continuous" or not packed_sequences:
            return tensors
        pos_ids = self._position_id.generate(packed_sequences)
        if pos_ids:
            tensors["position_ids"] = [torch.tensor(pos_ids, dtype=torch.int32)]
        return tensors
    def _to_tensors(self, keys: Dict[str, list]) -> Dict[str, List[torch.Tensor]]:
        """Convert packed per-key id lists to tensors.
        Honours ``config.output.dtype`` overrides per key; falls back to
        ``int32``.  Handles three shapes (see
        :func:`astrai.preprocessing.core.to_per_record_tensors` for the
        equivalent online-path helper):
        - ``List[int]`` per record → one tensor per record.
        - ``List[List[int]]`` per record (GRPO responses/masks) → one tensor
          per record, inner lists flattened.
        - ``List[int]`` for the whole shard (pre-packed keys) → single tensor.
        """
        tensors: Dict[str, List[torch.Tensor]] = {}
        for key, ids_list in keys.items():
            dt = _STR_TO_DTYPE.get(
                self.config.output.dtype.get(key, "int32"), torch.int32
            )
            if ids_list and isinstance(ids_list[0], list):
                tensors[key] = [
                    torch.tensor(
                        list(chain.from_iterable(ids))
                        if ids and isinstance(ids[0], list)
                        else ids,
                        dtype=dt,
                    )
                    for ids in ids_list
                ]
            else:
                tensors[key] = [
                    torch.tensor(list(chain.from_iterable(ids_list)), dtype=dt)
                ]
        return tensors
@@ -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,92 @@
 """Tokenization transform for JSONL record streams.
 Bridges the Reader layer (``JsonlStore`` reads raw JSON records) and the
 Dataset layer (expects per-record tensors).  Holds the tokenizer,
 mask-builder and position-id strategy together so that I/O code stays
 free of model dependencies.
 The record-processing core (mask building, primary-id extraction,
 per-key tensorisation, position-id generation) is shared with
 :class:`astrai.preprocessing.pipeline.Pipeline` via the
 :mod:`astrai.preprocessing.core` helpers.
 """
 import json
 from pathlib import Path
 from typing import Dict, List
 import torch
 from astrai.config.preprocess_config import PipelineConfig
 from astrai.preprocessing.core import (
    build_position_ids,
    build_preprocessing_components,
    iter_raw_records,
    to_per_record_tensors,
 )
 class TokenizeTransform:
    """Tokenize raw JSONL record dicts into per-key tensor lists.
    Owns the three preprocessing concerns that were previously inlined in
    ``JsonlStore``: tokenization, loss-mask construction and position-id
    generation.  Constructing it loads the tokenizer, so it is intentionally
    cheap to pass around once built.
    Args:
        config: Pipeline config describing sections / masks / position mode.
        tokenizer_path: Path passed to ``AutoTokenizer.from_pretrained``.
    """
    def __init__(self, config: PipelineConfig, tokenizer_path: str):
        self.config = config
        self.tokenizer, self.mask_builder, self.position_strategy = (
            build_preprocessing_components(config, tokenizer_path)
        )
    @classmethod
    def from_config_file(cls, config_path: str) -> "TokenizeTransform":
        """Build from a ``dataset_config.json`` file path.
        The config file follows :class:`PipelineConfig` schema with an
        extra ``tokenizer_path`` field.  When omitted, the config's
        parent directory is used as the tokenizer path.
        """
        root = Path(config_path).parent
        with open(config_path, "r", encoding="utf-8") as f:
            raw_config = json.load(f)
        tokenizer_path = raw_config.pop("tokenizer_path", None) or str(root)
        config = PipelineConfig.from_dict(raw_config)
        return cls(config, tokenizer_path)
    def apply(self, records: List[dict]) -> Dict[str, list]:
        """Tokenize a list of raw record dicts.
        Returns a dict mapping key (``sequence``, ``chosen``, ``responses``,
        …) to a list of per-record tensors (or nested tensor lists for
        multi-response keys such as GRPO ``responses``).
        """
        raw: Dict[str, list] = {}
        doc_sequences: List[List[int]] = []
        for result in iter_raw_records(
            records, self.mask_builder, self.config, self.tokenizer
        ):
            primary = None
            for val in result.values():
                if isinstance(val, list) and val and isinstance(val[0], int):
                    primary = val
                    break
            if primary is not None:
                doc_sequences.append(primary)
            for key, ids in result.items():
                raw.setdefault(key, []).append(ids)
        tensors = to_per_record_tensors(raw)
        pos_ids = build_position_ids(doc_sequences, self.position_strategy)
        if pos_ids is not None:
            tensors["position_ids"] = [torch.tensor(pos_ids, dtype=torch.int32)]
        return tensors
@@ -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,116 +0,0 @@
 import json
 import os
 from pathlib import Path
 from typing import Any, Dict, List, Optional
 import h5py
 import safetensors.torch as st
 import torch
 import torch.distributed as dist
 from torch import Tensor
 from astrai.parallel.setup import get_rank
 def save_h5(file_path: str, file_name: str, tensor_group: Dict[str, List[Tensor]]):
    os.makedirs(file_path, exist_ok=True)
    full_file_path = os.path.join(file_path, f"{file_name}.h5")
    with h5py.File(full_file_path, "w") as f:
        for key, tensors in tensor_group.items():
            grp = f.create_group(key)
            for idx, tensor in enumerate(tensors):
                arr = tensor.cpu().numpy()
                grp.create_dataset(f"data_{idx}", data=arr)
 def load_h5(file_path: str, share_memory=True) -> Dict[str, List[Tensor]]:
    tensor_group: Dict[str, List[Tensor]] = {}
    root_path = Path(file_path)
    h5_files = list(root_path.rglob("*.h5")) + list(root_path.rglob("*.hdf5"))
    for h5_file in h5_files:
        with h5py.File(h5_file, "r") as f:
            for key in f.keys():
                grp = f[key]
                dsets = []
                for dset_name in grp.keys():
                    dset = grp[dset_name]
                    tensor = torch.from_numpy(dset[:])
                    if share_memory:
                        tensor = tensor.share_memory_()
                    dsets.append(tensor)
                if tensor_group.get(key) is None:
                    tensor_group[key] = []
                tensor_group[key].extend(dsets)
    return tensor_group
 class Checkpoint:
    def __init__(
        self,
        state_dict: Dict[str, Any],
        epoch: int = 0,
        iteration: int = 0,
        extra: Optional[Dict[str, Any]] = None,
    ):
        self.state_dict = state_dict
        self.epoch = epoch
        self.iteration = iteration
        self.extra = extra or {}
    def save(
        self,
        save_dir: str,
    ) -> None:
        save_path = Path(save_dir)
        save_path.mkdir(parents=True, exist_ok=True)
        rank = get_rank()
        if rank == 0:
            meta = {
                "epoch": self.epoch,
                "iteration": self.iteration,
            }
            with open(save_path / "meta.json", "w") as f:
                json.dump(meta, f, indent=2)
            st.save_file(self.state_dict, save_path / "state_dict.safetensors")
            if self.extra:
                torch.save(self.extra, save_path / "extra.pt")
    @classmethod
    def load(
        cls,
        save_dir: str,
    ) -> "Checkpoint":
        rank = get_rank()
        save_path = Path(save_dir)
        meta = {}
        if rank == 0:
            with open(Path(save_dir) / "meta.json", "r") as f:
                meta = json.load(f)
        if dist.is_initialized():
            meta_list = [meta]
            dist.broadcast_object_list(meta_list, src=0)
            meta = meta_list[0]
        state_dict = st.load_file(save_path / "state_dict.safetensors")
        extra = None
        extra_path = save_path / "extra.pt"
        if extra_path.exists():
            extra = torch.load(extra_path, map_location="cpu", weights_only=False)
        return cls(
            state_dict=state_dict,
            epoch=meta["epoch"],
            iteration=meta["iteration"],
            extra=extra,
        )
@@ -0,0 +1,45 @@
 """Serialization utilities for models and datasets.
 This package re-exports checkpoint helpers and dataset storage helpers so
 that existing imports from ``astrai.serialization`` continue to work.
 """
 from astrai.serialization.checkpoint import (
    Checkpoint,
    load_json,
    load_model_config,
    load_model_weights,
    load_safetensors,
    load_state_dict,
    load_torch,
    save_json,
    save_model,
    save_safetensors,
    save_torch,
 )
 from astrai.serialization.dataset import (
    load_bin,
    load_bin_offsets,
    load_h5,
    save_bin,
    save_h5,
 )
 __all__ = [
    "Checkpoint",
    "load_json",
    "load_model_config",
    "load_model_weights",
    "load_safetensors",
    "load_state_dict",
    "load_torch",
    "save_json",
    "save_model",
    "save_safetensors",
    "save_torch",
    "load_bin",
    "load_bin_offsets",
    "load_h5",
    "save_bin",
    "save_h5",
 ]
@@ -0,0 +1,201 @@
 """Model checkpoint serialization helpers."""
 import io
 import json
 import time
 from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Any, Dict, Optional, Union
 import safetensors.torch as st
 import torch
 import torch.distributed as dist
 from astrai.parallel.setup import get_rank
 _META_FILE = "meta.json"
 _CONFIG_FILE = "config.json"
 _WEIGHTS_FILE = "model.safetensors"
 def save_safetensors(state_dict: dict, path: Union[str, Path]):
    st.save_file(state_dict, str(path))
 def load_safetensors(path: Union[str, Path], broadcast: bool = False) -> dict:
    if not broadcast or not dist.is_initialized():
        return st.load_file(str(path))
    rank = get_rank()
    if rank == 0:
        state_dict = st.load_file(str(path))
    else:
        state_dict = {}
    tmp = [state_dict]
    dist.broadcast_object_list(tmp, src=0)
    return tmp[0]
 def save_json(data: dict, path: Union[str, Path]):
    with open(str(path), "w") as f:
        json.dump(data, f, indent=2)
 def load_json(path: Union[str, Path], broadcast: bool = False) -> dict:
    if not broadcast or not dist.is_initialized():
        with open(str(path), "r") as f:
            return json.load(f)
    rank = get_rank()
    if rank == 0:
        with open(str(path), "r") as f:
            data = json.load(f)
    else:
        data = {}
    tmp = [data]
    dist.broadcast_object_list(tmp, src=0)
    return tmp[0]
 def save_torch(obj: Any, path: Union[str, Path]):
    torch.save(obj, str(path))
 def load_torch(path: Union[str, Path], broadcast: bool = False) -> Any:
    if not broadcast or not dist.is_initialized():
        return torch.load(str(path), map_location="cpu", weights_only=False)
    path = Path(path)
    rank = get_rank()
    if rank == 0:
        with open(path, "rb") as f:
            raw = f.read()
        data_tensor = torch.frombuffer(bytearray(raw), dtype=torch.uint8)
        num_bytes = torch.tensor([len(raw)], dtype=torch.long)
    else:
        num_bytes = torch.tensor([0], dtype=torch.long)
    dist.broadcast(num_bytes, src=0)
    if rank != 0:
        data_tensor = torch.empty(num_bytes.item(), dtype=torch.uint8)
    dist.broadcast(data_tensor, src=0)
    buf = io.BytesIO(data_tensor.numpy().tobytes())
    return torch.load(buf, map_location="cpu", weights_only=False)
 def save_model(config: dict, state_dict: dict, save_directory: str):
    save_path = Path(save_directory)
    save_path.mkdir(parents=True, exist_ok=True)
    save_json(config, save_path / _CONFIG_FILE)
    save_safetensors(state_dict, save_path / _WEIGHTS_FILE)
 def load_model_config(save_directory: str) -> dict:
    return load_json(Path(save_directory) / _CONFIG_FILE)
 def load_model_weights(save_directory: str) -> dict:
    return load_state_dict(Path(save_directory) / _WEIGHTS_FILE)
 def load_state_dict(path: Union[str, Path], broadcast: bool = False) -> dict:
    path = Path(path)
    if not broadcast or not dist.is_initialized():
        return load_safetensors(path)
    rank = get_rank()
    if rank == 0:
        state_dict = load_safetensors(path)
        specs = [
            (k, list(state_dict[k].shape), str(state_dict[k].dtype).split(".")[-1])
            for k in sorted(state_dict)
        ]
    else:
        state_dict = {}
        specs = []
    specs_list = [specs]
    dist.broadcast_object_list(specs_list, src=0)
    specs = specs_list[0]
    for key, shape, dtype_name in specs:
        dtype = getattr(torch, dtype_name)
        if rank != 0:
            tensor = torch.empty(shape, dtype=dtype, device="cpu")
        else:
            tensor = state_dict[key].contiguous().cpu()
        dist.broadcast(tensor, src=0)
        if rank != 0:
            state_dict[key] = tensor
    return state_dict
@dataclass
 class Checkpoint:
    state_dict: Dict[str, Any] = field(default_factory=dict)
    epoch: int = 0
    consumed_samples: int = 0
    extra: Dict[str, Any] = field(default_factory=dict)
    meta: Dict[str, Any] = field(default_factory=dict)
    config: Dict[str, Any] = field(default_factory=dict)
    def save(self, save_dir: str):
        save_path = Path(save_dir)
        save_path.mkdir(parents=True, exist_ok=True)
        meta = {
            "epoch": self.epoch,
            "consumed_samples": self.consumed_samples,
            "timestamp": time.strftime("%Y-%m-%dT%H:%M:%S"),
            **self.meta,
        }
        save_json(meta, save_path / _META_FILE)
        save_json(self.config, save_path / _CONFIG_FILE)
        save_safetensors(self.state_dict, save_path / _WEIGHTS_FILE)
        for key, value in self.extra.items():
            save_torch(value, save_path / f"{key}.pt")
    @classmethod
    def load(cls, save_dir: str, broadcast: bool = False) -> "Checkpoint":
        save_path = Path(save_dir)
        meta = load_json(save_path / _META_FILE, broadcast)
        config = load_json(save_path / _CONFIG_FILE, broadcast)
        state_dict = load_state_dict(save_path / _WEIGHTS_FILE, broadcast=broadcast)
        extra = {}
        for f in sorted(save_path.iterdir()):
            if f.suffix == ".pt":
                extra[f.stem] = load_torch(f, broadcast=broadcast)
        return cls(
            state_dict=state_dict,
            epoch=meta.get("epoch", 0),
            consumed_samples=meta.get("consumed_samples", 0),
            extra=extra,
            meta=meta,
            config=config,
        )
    @classmethod
    def load_any(cls, save_dir: str, broadcast: bool = False) -> Optional["Checkpoint"]:
        save_path = Path(save_dir)
        meta_path = save_path / _META_FILE
        weights_path = save_path / _WEIGHTS_FILE
        if meta_path.exists():
            return cls.load(save_dir, broadcast=broadcast)
        if weights_path.exists():
            state_dict = load_state_dict(weights_path, broadcast=broadcast)
            config = {}
            config_path = save_path / _CONFIG_FILE
            if config_path.exists():
                config = load_json(config_path, broadcast)
            return cls(state_dict=state_dict, config=config)
        return None
@@ -0,0 +1,123 @@
 """Dataset storage serialization helpers (HDF5 / memory-mapped binary)."""
 import json
 import os
 from pathlib import Path
 from typing import Any, Dict, List, Optional
 import h5py
 import numpy as np
 import torch
 from torch import Tensor
 def save_h5(file_path: str, file_name: str, tensor_group: Dict[str, List[Tensor]]):
    os.makedirs(file_path, exist_ok=True)
    full_file_path = os.path.join(file_path, f"{file_name}.h5")
    with h5py.File(full_file_path, "w") as f:
        for key, tensors in tensor_group.items():
            grp = f.create_group(key)
            for idx, tensor in enumerate(tensors):
                arr = tensor.cpu().numpy()
                grp.create_dataset(f"data_{idx}", data=arr)
 def load_h5(file_path: str, share_memory=True) -> Dict[str, List[Tensor]]:
    tensor_group: Dict[str, List[Tensor]] = {}
    root_path = Path(file_path)
    if root_path.is_file() and root_path.suffix in (".h5", ".hdf5"):
        h5_files = [root_path]
    else:
        h5_files = list(root_path.rglob("*.h5")) + list(root_path.rglob("*.hdf5"))
    for h5_file in h5_files:
        with h5py.File(h5_file, "r") as f:
            for key in f.keys():
                grp = f[key]
                dsets = []
                for dset_name in grp.keys():
                    dset = grp[dset_name]
                    tensor = torch.from_numpy(dset[:])
                    if share_memory:
                        tensor = tensor.share_memory_()
                    dsets.append(tensor)
                if tensor_group.get(key) is None:
                    tensor_group[key] = []
                tensor_group[key].extend(dsets)
    return tensor_group
 def save_bin(
    file_path: str,
    tensor_group: Dict[str, List[Tensor]],
    record_keys: Optional[List[str]] = None,
 ):
    """Save tensors as memory-mapped binary files.
    When *record_keys* is provided, those keys are written with per-record
    cumulative offsets in ``meta.json`` so that ``MmapStore.fetch_record``
    can slice individual records from the concatenated binary without
    cross-record concatenation.  Keys not in *record_keys* (e.g. SEQ
    ``sequence``) are written as a single contiguous stream without
    offsets, preserving backward compatibility.
    Nested keys (``List[List[Tensor]]`` such as GRPO ``responses``) are
    not supported in bin format — use H5 for those.
    """
    os.makedirs(file_path, exist_ok=True)
    record_keys = set(record_keys or [])
    meta = {}
    for key, tensors in tensor_group.items():
        if tensors and isinstance(tensors[0], list):
            raise ValueError(
                f"Nested key '{key}' (List[List[Tensor]]) is not supported "
                f"in bin format. Use H5 or JSONL storage instead."
            )
        cat = torch.cat(tensors, dim=0)
        entry: Dict[str, Any] = {
            "shape": list(cat.shape),
            "dtype": str(cat.dtype).split(".")[-1],
        }
        if key in record_keys:
            offsets = [0]
            for t in tensors:
                offsets.append(offsets[-1] + t.shape[0])
            entry["offsets"] = offsets
        meta[key] = entry
        np.asarray(cat.cpu().numpy()).tofile(os.path.join(file_path, f"{key}.bin"))
    with open(os.path.join(file_path, "meta.json"), "w") as f:
        json.dump(meta, f)
 def load_bin(file_path: str) -> Dict[str, List[Tensor]]:
    with open(os.path.join(file_path, "meta.json"), "r") as f:
        meta = json.load(f)
    segments: Dict[str, List[Tensor]] = {}
    for key, info in meta.items():
        arr = np.memmap(
            os.path.join(file_path, f"{key}.bin"),
            dtype=info["dtype"],
            mode="r",
            shape=tuple(info["shape"]),
        )
        segments[key] = [torch.from_numpy(arr)]
    return segments
 def load_bin_offsets(file_path: str) -> Dict[str, List[int]]:
    """Read per-record cumulative offsets from ``meta.json``.
    Returns an empty dict when no key has offsets (legacy bin files),
    in which case record-mode access falls back to per-record segment
    indexing (H5/JSONL layout).
    """
    with open(os.path.join(file_path, "meta.json"), "r") as f:
        meta = json.load(f)
    offsets: Dict[str, List[int]] = {}
    for key, info in meta.items():
        if "offsets" in info:
            offsets[key] = info["offsets"]
    return offsets
@@ -1,13 +1,11 @@
-from dataclasses import dataclass
+from functools import cached_property
 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 +13,36 @@ 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,
        special_tokens: Optional[Dict[str, str]] = None,
    ):
        self.name = name
        self.template_str = template_str
        self.description = description
        self.default_variables = default_variables or {}
        self.special_tokens = special_tokens or {}
-    def __post_init__(self):
+    @cached_property
-        if self.default_variables is None:
+    def _compiled(self) -> Template:
-            self.default_variables = {}
+        """Lazy-compiled Jinja2 template, cached on first access.
-        if self.special_tokens is None:
+
-            self.special_tokens = {}
+        The compiled :class:`~jinja2.Template` holds a dynamically-generated
        ``root`` render function whose ``__module__`` is ``None``; under
        ``pickle`` it falls back to ``__main__`` and breaks ``spawn``-based
        multiprocessing.  By deferring compilation to first access, the
        default pickle protocol serialises only ``template_str``; each
        worker rebuilds the cache on first render.
        """
        return Template(self.template_str)
    @classmethod
    def from_string(
@@ -43,7 +54,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 +84,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):
@@ -147,7 +164,14 @@ class AutoTokenizer:
          - tokenizer.bos_token   → returns string
          - tokenizer.bos_token_id → returns corresponding integer ID
          - tokenizer.stop_ids → returns list of corresponding integer IDs for all special tokens
        Internal/private attrs are not intercepted: during unpickle
        ``__dict__`` is empty, so probing ``self._special_token_map``
        would recurse infinitely.
        """
        if key.startswith("_"):
            raise AttributeError(key)
        # Handle stop_ids - return IDs for all special tokens
        if key == "stop_ids":
            stop_ids = []
@@ -1,75 +1,25 @@
 from typing import Dict
 import torch
 import torch.nn as nn
-def grad_norm(model: nn.Module, norm_type: int = 2) -> Dict[str, float]:
+def grad_norm(model: nn.Module, per_param: bool = False) -> float | Dict[str, float]:
-    """Compute gradient norm for each parameter in the model."""
+    grads = [p.grad.detach() for p in model.parameters() if p.grad is not None]
    if not grads:
        return 0.0
    total_sq = torch.stack([g.pow(2).sum() for g in grads]).sum()
    if per_param:
        norms = {}
        for name, param in model.named_parameters():
            if param.grad is not None:
                norms[name] = param.grad.norm(2).item()
            else:
                norms[name] = 0.0
-        if param.grad:
+        norms["total"] = total_sq.sqrt().item()
            norm = param.grad.data.norm(norm_type).item()
            norms[name] = norm
        return norms
-
+    return total_sq.sqrt().item()
 def grad_std(model: nn.Module) -> Dict[str, float]:
    """Compute standard deviation of gradients for each parameter."""
    stds = {}
    for name, param in model.named_parameters():
        stds[name] = 0.0
        if param.grad:
            std = param.grad.data.std().item()
            stds[name] = std
    return stds
 def grad_max(model: nn.Module) -> Dict[str, float]:
    """Find the maximum absolute gradient value for each parameter."""
    max_vals = {}
    for name, param in model.named_parameters():
        max_vals[name] = -float("inf")
        if param.grad:
            max_val = param.grad.data.max().item()
            max_vals[name] = max_val
    return max_vals
 def grad_min(model: nn.Module) -> Dict[str, float]:
    """Find the minimum absolute gradient value for each parameter."""
    min_vals = {}
    for name, param in model.named_parameters():
        min_vals[name] = float("inf")
        if param.grad:
            min_val = param.grad.data.min().item()
            min_vals[name] = min_val
    return min_vals
 def grad_mean(model: nn.Module) -> Dict[str, float]:
    """Compute mean of gradients for each parameter."""
    means = {}
    for name, param in model.named_parameters():
        means[name] = 0.0
        if param.grad:
            mean = param.grad.data.mean().item()
            means[name] = mean
    return means
 def grad_nan_num(model: nn.Module) -> Dict[str, int]:
    """Count the number of NaNs in gradients for each parameter."""
    nan_nums = {}
    for name, param in model.named_parameters():
        nan_nums[name] = 0
        if param.grad:
            nan_num = param.grad.isnan().sum().item()
            nan_nums[name] = nan_num
    return nan_nums
 def ctx_get_loss(ctx):
@@ -80,25 +30,9 @@ def ctx_get_lr(ctx):
    return ctx.optimizer.param_groups[-1]["lr"]
 def ctx_get_val_loss(ctx):
    return ctx.val_loss
 def ctx_get_grad_norm(ctx):
-    return grad_norm(ctx.model)
+    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,34 @@ 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.extra_kwargs = kwargs
    @abstractmethod
@@ -122,32 +131,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 +143,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 +174,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(),
@@ -223,12 +222,13 @@ class DPOStrategy(BaseStrategy):
        self,
        model: nn.Module,
        device: str,
        ref_model: nn.Module,
        beta: float = 0.1,
-        reduction: str = "mean",
+        reduction: str = "sum",
        **kwargs,
    ):
        super().__init__(model, device, **kwargs)
-        self.ref_model = create_ref_model(model)
+        self.ref_model = ref_model
        self.beta = beta
        self.reduction = reduction
@@ -265,41 +265,45 @@ class DPOStrategy(BaseStrategy):
 class GRPOStrategy(BaseStrategy):
    """Group Relative Policy Optimization strategy.
-    On-policy GRPO following DeepSeek-R1: the policy model is updated while
+    Implements GRPO following DeepSeek-R1 with token-level PPO clipping.
-    a frozen ref_model stores the old-policy log-probs.  ratio = exp(logπ_θ - logπ_ref),
+    Advantages are group-normalized from scalar per-response rewards and
-    clipped PPO objective.  Call ``sync_ref_model()`` after each data-generation round.
+    broadcast across all response tokens.  The loss is computed **only on
    response tokens** — prompt tokens are masked out.
    Three model roles are distinguished:
    * **Policy** ``self.model`` — the model being trained.
    * **Old policy** ``self.old_model`` — the behaviour policy that generated
      the responses.  Used for the importance sampling ratio
      ``ρ = π_θ / π_old``.  Synced externally after each data-generation round.
    * **Reference model** ``self.ref_model`` — a frozen copy of the initial
      policy (typically the SFT checkpoint) used **only** for the KL
      regularisation term.  It is never updated during training.
    """
    def __init__(
        self,
        model: nn.Module,
        device: str,
        old_model: nn.Module,
        ref_model: nn.Module,
        clip_eps: float = 0.2,
        kl_coef: float = 0.01,
        group_size: int = 4,
        reduction: str = "mean",
        sync_interval: int = 200,
        **kwargs,
    ):
        super().__init__(model, device, **kwargs)
-        self.ref_model = create_ref_model(model)
+        self.old_model = old_model
        self.ref_model = ref_model
        self.clip_eps = clip_eps
        self.kl_coef = kl_coef
        self.group_size = group_size
        self.reduction = reduction
        self.sync_interval = sync_interval
        self._step = 0
-    def sync_ref_model(self):
+    def sync_old_model(self):
-        """Copy current model weights to ref model."""
+        """Copy current policy weights to old model."""
-        ref_state = self.model.state_dict()
+        self.old_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
        if self._step % self.sync_interval == 0:
            self.sync_ref_model()
        batch = move_to_device(batch, self.device)
        prompts = batch["prompts"]
        responses = batch["responses"]
@@ -310,33 +314,60 @@ class GRPOStrategy(BaseStrategy):
        responses_flat = responses.view(-1, response_len)
        masks_flat = masks.view(-1, response_len)
        prompt_expanded = prompts.unsqueeze(1).repeat(1, group_size, 1).flatten(0, 1)
        prompt_len = prompt_expanded.size(1)
        full_sequences = torch.cat([prompt_expanded, responses_flat], dim=-1)
-        full_masks = torch.cat([torch.ones_like(prompt_expanded), masks_flat], dim=-1)
+        # Prompt tokens are masked out (0) so logprobs are computed only for
-
+        # response tokens.  get_logprobs shifts the mask by one position, so
-        log_probs_policy = get_logprobs(
+        # the first response token's logprob (predicted from the last prompt
-            self.model, full_sequences, full_masks, self.reduction
+        # token) is correctly included.
-        )
+        full_masks = torch.cat([torch.zeros_like(prompt_expanded), masks_flat], dim=-1)
        log_probs_policy = log_probs_policy.view(batch_size, group_size)
        # get_logprobs returns [B*G, S-1] (S = prompt_len + response_len).
        # Response token logprobs occupy the last ``response_len`` positions
        # (the first response token is predicted from the last prompt token).
        token_log_probs_policy = get_logprobs(
            self.model, full_sequences, full_masks, "none"
        )[:, prompt_len - 1 :]
        with torch.no_grad():
-            log_probs_ref = get_logprobs(
+            token_log_probs_old = get_logprobs(
-                self.ref_model, full_sequences, full_masks, self.reduction
+                self.old_model, full_sequences, full_masks, "none"
-            )
+            )[:, prompt_len - 1 :]
-            log_probs_ref = log_probs_ref.view(batch_size, group_size)
+            token_log_probs_ref = get_logprobs(
                self.ref_model, full_sequences, full_masks, "none"
            )[:, prompt_len - 1 :]
-        eps = torch.finfo(log_probs_policy.dtype).eps
+        # Reshape to [B, G, response_len]
        token_log_probs_policy = token_log_probs_policy.view(batch_size, group_size, -1)
        token_log_probs_old = token_log_probs_old.view(batch_size, group_size, -1)
        token_log_probs_ref = token_log_probs_ref.view(batch_size, group_size, -1)
        token_masks = masks_flat.view(batch_size, group_size, -1).float()
        # Group-normalized advantages from scalar per-response rewards.
        eps = 1e-8
        mean = rewards.mean(dim=-1, keepdim=True)
-        std = rewards.std(dim=-1, keepdim=True)
+        std = rewards.std(dim=-1, keepdim=True, unbiased=False)
        advantages = (rewards - mean) / (std + eps)
        # Broadcast scalar advantage to every response token: [B, G, 1]
        advantages = advantages.unsqueeze(-1)
-        ratio = torch.exp(log_probs_policy - log_probs_ref)
+        # Token-level ratio (π_θ / π_old) and PPO clipping.
        log_ratio = token_log_probs_policy - token_log_probs_old
        ratio = torch.exp(log_ratio)
        surr1 = ratio * advantages
        surr2 = torch.clamp(ratio, 1 - self.clip_eps, 1 + self.clip_eps) * advantages
        per_token_policy_loss = -torch.min(surr1, surr2)
        token_count = token_masks.sum().clamp(min=1.0)
        policy_loss = (per_token_policy_loss * token_masks).sum() / token_count
        # KL penalty to frozen reference model with k1 estimator (non-negative):
        # k1 = π_ref / π_θ - log(π_ref / π_θ) - 1, where π_ref / π_θ = exp(log_ref - log_policy).
        log_ref_ratio = token_log_probs_ref - token_log_probs_policy
        r = torch.exp(log_ref_ratio)
        kl_per_token = r - torch.log(r + eps) - 1.0
        kl_penalty = self.kl_coef * (kl_per_token * token_masks).sum() / token_count
        policy_loss = -torch.min(surr1, surr2).mean()
        kl_penalty = self.kl_coef * (log_probs_policy - log_probs_ref).square().mean()
        total_loss = policy_loss + kl_penalty
        return total_loss
@@ -1,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
 from astrai.serialization import Checkpoint
 from astrai.trainer.metric_util import (
    ctx_get_grad_max,
    ctx_get_grad_mean,
    ctx_get_grad_min,
    ctx_get_grad_nan_num,
    ctx_get_grad_norm,
    ctx_get_grad_std,
    ctx_get_loss,
    ctx_get_lr,
    ctx_get_val_loss,
 )
 from astrai.trainer.train_context import TrainContext
 logger = logging.getLogger(__name__)
@runtime_checkable
 class TrainCallback(Protocol):
@@ -42,18 +45,15 @@ class TrainCallback(Protocol):
    def on_epoch_end(self, context: TrainContext):
        """Called at the end of each epoch."""
    def on_step_begin(self, context: TrainContext):
        """Called at the beginning of each step."""
    def on_step_end(self, context: TrainContext):
        """Called at the end of each step."""
    def on_batch_begin(self, context: TrainContext):
        """Called at the beginning of each batch."""
    def on_batch_end(self, context: TrainContext):
        """Called at the end of each batch."""
    def on_optimizer_step(self, context: TrainContext):
        """Called on every optimizer step (sync step only)."""
    def on_error(self, context: TrainContext):
        """Called when an error occurs during training."""
@@ -69,12 +69,6 @@ class CallbackFactory(BaseFactory[TrainCallback]):
        callback = CallbackFactory.create("my_callback", **kwargs)
    """
    @classmethod
    def _validate_component(cls, callback_cls: type) -> None:
        """Validate that the callback class inherits from TrainCallback."""
        if not issubclass(callback_cls, TrainCallback):
            raise TypeError(f"{callback_cls.__name__} must inherit from TrainCallback")
@CallbackFactory.register("gradient_clipping")
 class GradientClippingCallback(TrainCallback):
@@ -85,28 +79,45 @@ class GradientClippingCallback(TrainCallback):
    def __init__(self, max_grad_norm: float):
        self.max_grad_norm = max_grad_norm
-    def on_step_begin(self, context: TrainContext):
+    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("scheduler")
+@CallbackFactory.register("gradient_checkpointing")
-class SchedulerCallback(TrainCallback):
+class GradientCheckpointingCallback(TrainCallback):
    """
-    Scheduler callback for trainer.
+    Activation checkpointing callback — trades compute for memory
    by recomputing specified module activations during the backward pass.
    Args:
        modules: Module types to apply checkpointing to.
    """
-    def __init__(self):
+    def __init__(self, modules: Optional[List[type]] = None):
-        pass
+        self.modules = tuple(modules) if modules else ()
    def _enable(self, module: nn.Module):
        if self.modules and isinstance(module, self.modules):
            fn = module.forward
            module._original_forward = fn
            module.forward = lambda *a, **kw: torch_checkpoint(
                fn, *a, use_reentrant=False, **kw
            )
    @staticmethod
    def _disable(module: nn.Module):
        if hasattr(module, "_original_forward"):
            module.forward = module._original_forward
            del module._original_forward
    def on_train_begin(self, context: TrainContext):
-        for group in context.optimizer.param_groups:
+        context.model.apply(self._enable)
-            if "initial_lr" not in group:
+        logger.info("Gradient checkpointing enabled")
                group["initial_lr"] = group["lr"]
-    def on_batch_end(self, context: TrainContext):
+    def on_train_end(self, context: TrainContext):
-        if context.scheduler:
+        context.model.apply(self._disable)
            context.scheduler.step()
@CallbackFactory.register("checkpoint")
@@ -115,54 +126,65 @@ 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 = None
-        self.last_ckpt_iter = 0
+
    def on_train_begin(self, context: TrainContext):
        self.last_ckpt_step = context.optimizer_step
    @only_on_rank(0)
    def _save_checkpoint(self, context: TrainContext):
-        save_path = os.path.join(
+        self.last_ckpt_step = context.optimizer_step
            self.save_dir, f"epoch_{context.epoch}_iter_{context.iteration}"
        )
        state_dict = (
            self.state_dict_fn(context.model)
            if self.state_dict_fn
            else context.model.state_dict()
        )
-        extra = self.save_extra_fn(context) if self.save_extra_fn else None
+        with context.executor.checkpoint_context(context.model) as state_dict:
            if state_dict is not None:
                save_path = os.path.join(
                    self.save_dir,
                    f"epoch_{context.epoch}_step_{context.optimizer_step}",
                )
                extra = self.save_extra_fn(context)
                meta = context.config.to_dict()
                context.checkpoint = Checkpoint(
                    state_dict=state_dict,
                    epoch=context.epoch,
-            iteration=context.iteration,
+                    consumed_samples=context.consumed_samples,
                    config=context.model_config,
                    extra=extra,
                    meta=meta,
                )
                context.checkpoint.save(save_path)
        self.last_ckpt_iter = context.iteration
    def on_batch_end(self, context: TrainContext):
-        if context.iteration - self.last_ckpt_iter >= self.interval:
+        if context.optimizer_step - self.last_ckpt_step >= self.interval:
            self._save_checkpoint(context)
    def on_train_end(self, context: TrainContext):
-        if context.iteration != self.last_ckpt_iter:
+        if context.optimizer_step != self.last_ckpt_step:
            self._save_checkpoint(context)
    def on_error(self, context: TrainContext):
        self._save_checkpoint(context)
    @staticmethod
    def save_extra(context: TrainContext) -> dict:
        extra = {}
        for name in CheckpointCallback.extra_keys:
            obj = getattr(context, name, None)
            if obj:
                extra[name] = obj.state_dict()
        return extra
@CallbackFactory.register("progress_bar")
 class ProgressBarCallback(TrainCallback):
@@ -170,26 +192,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": f"{context.optimizer_step:d}",
            "loss": f"{context.loss:.4f}",
            "lr": f"{context.optimizer.param_groups[-1]['lr']:.2e}",
        }
-        )
+        if context.grad_norm is not None:
            postfix["grad_norm"] = f"{context.grad_norm:.2f}"
        if context.val_loss is not None:
            postfix["val_loss"] = f"{context.val_loss:.4f}"
        self.progress_bar.set_postfix(postfix)
        self.progress_bar.update(1)
    @only_on_rank(0)
@@ -199,19 +231,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 = None
        self.save_interval = save_interval
        self.log_interval = log_interval
        self.metrics = metrics or ["loss", "lr"]
        self.val_step = val_step
        self._next_val_step = 0
        self.log_dir = Path(log_dir) if log_dir else Path.cwd() / "logs"
        self.log_dir.mkdir(parents=True, exist_ok=True)
@@ -221,46 +254,91 @@ 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
    def on_train_begin(self, context: TrainContext):
        self.last_log_flush_step = context.optimizer_step
    @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 (
-            self._save_log(context.epoch, context.iteration)
+            self.last_log_flush_step is None
            or context.optimizer_step != self.last_log_flush_step
        ):
            self._flush(context.epoch, context.optimizer_step)
            self.last_log_flush_step = context.optimizer_step
    def on_error(self, context):
-        self._save_log(context.epoch, context.iteration)
+        self._flush(context.epoch, context.optimizer_step)
@@ -1,16 +1,19 @@
 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.dataset import RDSampler
 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.trainer.strategy import BaseStrategy, StrategyFactory
+from astrai.serialization import Checkpoint, load_json
 from astrai.trainer.strategy import BaseStrategy, StrategyFactory, create_ref_model
@dataclass
@@ -18,84 +21,197 @@ 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._param_path: Optional[str] = None
-        self._load_extra_fn = load_extra_fn
+        self._resume: bool = False
-    def with_checkpoint(self, checkpoint: Optional[Checkpoint]) -> Self:
+    def with_param_path(self, param_path: Optional[str], resume: bool = False) -> Self:
-        self._checkpoint = checkpoint
+        self._param_path = param_path
        self._resume = resume
        return self
    def build(self) -> TrainContext:
        cfg = self.config
        device = get_current_device()
        executor = ExecutorFactory.create(
            cfg.parallel_mode,
            grad_accum_steps=cfg.grad_accum_steps,
            **cfg.executor_kwargs,
        )
        model = cfg.model_fn()
        model = model.to(device=device)
        model_config = {}
        if self._param_path:
            config_path = Path(self._param_path) / "config.json"
            if config_path.exists():
                model_config = load_json(config_path)
        if not model_config and hasattr(model, "config"):
            model_config = model.config.to_dict()
        context = TrainContext(
-            model=self.config.model,
+            model=model,
            world_size=get_world_size(),
            rank=get_rank(),
            config=cfg,
            model_config=model_config,
            executor=executor,
        )
-        device = get_current_device()
+        if self._param_path:
-        context.model = context.model.to(device=device)
+            checkpoint = Checkpoint.load_any(self._param_path)
            if checkpoint is not None:
                model.load_state_dict(checkpoint.state_dict, strict=False)
                if checkpoint.config:
                    context.model_config = checkpoint.config
-        if self.config.nprocs > 1 and self.config.parallel_wrapper:
+                if self._resume:
-            context.model = self.config.parallel_wrapper(context.model)
+                    context.epoch = checkpoint.epoch or cfg.start_epoch
-
+                    if checkpoint.consumed_samples > 0:
-        if self._checkpoint is not None:
+                        per_step = (
-            context.epoch = max(self._checkpoint.epoch, self.config.start_epoch)
+                            cfg.batch_per_device
-            context.iteration = max(self._checkpoint.iteration, self.config.start_batch)
+                            * context.world_size
-            context.model.load_state_dict(self._checkpoint.state_dict)
+                            * cfg.grad_accum_steps
-            context.checkpoint = self._checkpoint
+                        )
                        context.consumed_samples = (
                            checkpoint.consumed_samples // per_step
                        ) * per_step
                    else:
-            context.checkpoint = Checkpoint(
+                        context.consumed_samples = (
-                state_dict=context.model.state_dict(),
+                            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)
-        sampler = ResumableDistributedSampler(
+            n_val = max(1, int(n_total * cfg.val_split))
-            data_source=cfg.dataset,
+            n_train = n_total - n_val
            generator = torch.Generator().manual_seed(cfg.random_seed)
            train_dataset, val_dataset = random_split(
                cfg.dataset, [n_train, n_val], generator=generator
            )
        sampler_offset = context.consumed_samples // context.world_size
        sampler = RDSampler(
            data_source=train_dataset,
            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,
            collate_fn=cfg.collate_fn,
        )
        if val_dataset is not None:
            val_sampler = RDSampler(
                data_source=val_dataset,
                start_epoch=0,
                start_iter=0,
                seed=cfg.random_seed,
                shuffle=False,
            )
            context.val_dataloader = DataLoader(
                val_dataset,
                batch_size=cfg.batch_per_device,
                sampler=val_sampler,
                num_workers=cfg.num_workers,
                pin_memory=cfg.pin_memory,
                prefetch_factor=cfg.prefetch_factor,
                collate_fn=cfg.collate_fn,
            )
        context.model, context.optimizer, context.dataloader, context.scheduler = (
            executor.prepare(
                model,
                context.optimizer,
                context.dataloader,
                context.scheduler,
            )
        )
        if context.checkpoint and context.checkpoint.extra:
            extra = context.checkpoint.extra
            for name in ("optimizer", "scheduler"):
                if name in extra:
                    obj = getattr(context, name, None)
                    if obj is not None:
                        obj.load_state_dict(extra[name])
        strategy_kwargs = dict(cfg.extra_kwargs)
        if cfg.strategy in ("dpo", "grpo"):
            ref_model = create_ref_model(
                cfg.model_fn, executor.unwrap_model(context.model)
            ).to(device=device)
            strategy_kwargs["ref_model"] = ref_model
        if cfg.strategy == "grpo":
            old_model = create_ref_model(
                cfg.model_fn, executor.unwrap_model(context.model)
            ).to(device=device)
            strategy_kwargs["old_model"] = old_model
        context.strategy = StrategyFactory.create(
            cfg.strategy,
            model=context.model,
            train_type=self.config.strategy,
            device=device,
-            **self.config.extra_kwargs,
+            executor=executor,
            **strategy_kwargs,
        )
        return context
@@ -3,7 +3,6 @@ from typing import List, Optional
 from astrai.config import TrainConfig
 from astrai.parallel.setup import spawn_parallel_fn
 from astrai.serialization import Checkpoint
 from astrai.trainer.train_callback import (
    CallbackFactory,
    TrainCallback,
@@ -25,18 +24,27 @@ class Trainer:
    def _get_default_callbacks(self) -> List[TrainCallback]:
        cfg = self.train_config
-        return [
+        callbacks = [
            CallbackFactory.create(
                "gradient_checkpointing",
                modules=cfg.gradient_checkpointing_modules,
            ),
            CallbackFactory.create(
                "checkpoint",
                cfg.ckpt_dir,
                cfg.ckpt_interval,
            ),
            CallbackFactory.create(
                "metric",
                log_dir=cfg.log_dir,
                save_interval=cfg.ckpt_interval,
                metrics=cfg.metrics,
                val_step=cfg.val_step,
            ),
            CallbackFactory.create("progress_bar", cfg.n_epoch),
            CallbackFactory.create("checkpoint", cfg.ckpt_dir, cfg.ckpt_interval),
            CallbackFactory.create("metric_logger", cfg.ckpt_dir, cfg.ckpt_interval),
            CallbackFactory.create("gradient_clipping", cfg.max_grad_norm),
            CallbackFactory.create("scheduler"),
        ]
-
+        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,55 +52,61 @@ class Trainer:
            if method:
                method(context)
-    def train(self, checkpoint: Optional[Checkpoint] = None):
+    def _trainer_loop(self, param_path: Optional[str] = None, resume: bool = False):
-        config = self.train_config
+        context = (
-        spawn_parallel_fn(
+            TrainContextBuilder(self.train_config)
-            self._train_impl,
+            .with_param_path(param_path, resume=resume)
-            backend=config.backend,
+            .build()
            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()
-            # 1.epoch
+
-            for epoch in range(context.epoch, self.train_config.n_epoch):
+            for epoch in range(context.epoch, context.config.n_epoch):
                context.epoch = epoch
                self._call_callbacks("on_epoch_begin", context)
                accumulation_steps = max(self.train_config.accumulation_steps, 1)
                for batch in context.dataloader:
-                    if context.iteration % accumulation_steps == 0:
+                    with executor.accumulate(context.model):
                        # 2. step
                        self._call_callbacks("on_step_begin", context)
                        context.optimizer.step()
                        context.optimizer.zero_grad()
                        self._call_callbacks("on_step_end", context)
                    # 3. batch
                        self._call_callbacks("on_batch_begin", context)
                        loss = context.strategy(batch)
                        context.loss = loss.item()
-                    context.iteration += 1
+                        stand_loss = loss / executor.grad_accum_steps
-
+                        executor.backward(stand_loss)
-                    # to make the loss normalized by accumulation steps
+                        context.consumed_samples += (
-                    stand_loss = loss / accumulation_steps
+                            context.config.batch_per_device * context.world_size
-                    stand_loss.backward()
+                        )
                        self._call_callbacks("on_batch_end", context)
                        if executor.sync_gradients:
                            self._call_callbacks("on_optimizer_step", context)
                            context.optimizer.step()
                            context.optimizer.zero_grad()
                            if context.scheduler:
                                context.scheduler.step()
                self._call_callbacks("on_epoch_end", context)
        except Exception as e:
-            logger.error(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, param_path: Optional[str] = None, resume: bool = False):
        cfg = self.train_config
        spawn_parallel_fn(
            self._trainer_loop,
            backend=cfg.backend,
            world_size=cfg.nprocs,
            master_addr=cfg.master_addr,
            master_port=cfg.master_port,
            device_type=cfg.device_type,
            start_method=cfg.start_method,
            param_path=param_path,
            resume=resume,
        )
@@ -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,68 @@
 #pragma once
 template<typename T, typename AT = float>
 struct AttentionParams {
    int batch;
    int q_head;
    int kv_head;
    int q_len;
    int kv_len;
    int head_dim;
    int use_mask;
    int causal_offset;   // -1 = non-causal; >=0 = absolute position of first Q token
    int num_splits;
    float scale;
    // Q strides (element offsets for each dim — layout-agnostic)
    int q_stride_b, q_stride_h, q_stride_l, q_stride_d;
    // KV strides (K and V share the same layout — only base pointers differ)
    int kv_stride_b, kv_stride_h, kv_stride_l, kv_stride_d;
    // Mask: 2D [batch, kv_len] (mask_q_stride=0) or 3D [batch, q_len, kv_len]
    int mask_b_stride;   // = kv_len (both 2D and 3D)
    int mask_q_stride;   // 2D: 0 (all q rows share); 3D: kv_len
    const T* __restrict__ q;
    const T* __restrict__ k;
    const T* __restrict__ v;
    const bool* __restrict__ mask;
    T* __restrict__ o;
    AT* __restrict__ o_part;
    AT* __restrict__ ml_part;
 };
 template<typename T, typename AT = float>
 struct PagedAttentionParams {
    int batch;
    int q_head;
    int kv_head;
    int q_len;
    int kv_len;
    int head_dim;
    int use_mask;
    int causal_offset;
    float scale;
    int num_splits;
    int page_size;
    int max_pages;
    // Q strides (layout-agnostic)
    int q_stride_b, q_stride_h, q_stride_l, q_stride_d;
    // Mask strides (2D or 3D)
    int mask_b_stride;
    int mask_q_stride;
    const T* __restrict__ q;
    const T* __restrict__ k_cache;
    const T* __restrict__ v_cache;
    const bool* __restrict__ mask;
    const int64_t* __restrict__ page_table;
    T* __restrict__ o;
    AT* __restrict__ o_part;
    AT* __restrict__ ml_part;
 };
@@ -0,0 +1,82 @@
 #include "attn_decode_split_kv.cuh"
 #include "attn_entry_utils.cuh"
 #ifndef ASTRAI_NO_MMA
 #include "attn_decode_split_kv_mma.cuh"
 #endif
 // Scalar fallback: one warp per query head, split-KV across grid.z.
 static void launch_scalar_decode(AttentionParams<bf16>& p) {
    int group_size = p.q_head / p.kv_head;
    int chunks_total = (p.kv_len + DC_CHUNK - 1) / DC_CHUNK;
    p.num_splits = compute_num_splits(p.batch * p.kv_head, chunks_total);
    alloc_split_partials(p);
    size_t smem = DC_CHUNK * p.head_dim * sizeof(bf16);
    attn_decode_split_kv_kernel<<<dim3(p.batch * p.kv_head, 1, p.num_splits), dim3(32, group_size), smem>>>(p);
    attn_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim>>>(p);
 }
 #ifndef ASTRAI_NO_MMA
 // MMA head-packing requires G <= 16 (BR=16 rows). sm_80+ tensor-core
 // + cp.async wins even at G=1 (decode is memory-bound, not compute-bound).
 // STAGES=2 (double-buffer) for D<=128 (smem 16 KB); STAGES=1 for D=256
 // (double-buffer would be 32 KB, near the 48 KB static cap — keep single
 // to preserve occupancy).
 template <int HEAD_DIM, int BC, int STAGES = (HEAD_DIM <= 128) ? 2 : 1>
 static void launch_mma_decode(AttentionParams<bf16>& p) {
    int tiles_total = (p.kv_len + BC - 1) / BC;
    p.num_splits = compute_num_splits(p.batch * p.kv_head, tiles_total);
    alloc_split_partials(p);
    attn_decode_split_kv_mma_kernel<HEAD_DIM, BC, STAGES><<<dim3(p.kv_head, p.batch, p.num_splits), 32>>>(p);
    attn_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim>>>(p);
 }
 #endif
 template <int HEAD_DIM>
 static void dispatch_decode(AttentionParams<bf16>& p) {
 #ifndef ASTRAI_NO_MMA
    int G = p.q_head / p.kv_head;
    if (G >= 1 && G <= 16) {
        launch_mma_decode<HEAD_DIM, 32>(p);
        return;
    }
 #endif
    launch_scalar_decode(p);
 }
 torch::Tensor attn_decode(
    torch::Tensor q,
    torch::Tensor k,
    torch::Tensor v,
    c10::optional<torch::Tensor> mask,
    int64_t causal_offset,
    double scale,
    int64_t layout
 ) {
    AttentionParams<bf16> p;
    attn_pack_params(q, k, v, mask, causal_offset, scale, layout, p);
    TORCH_CHECK(p.q_len == 1, "Q seq_len must be 1");
    TORCH_CHECK(p.head_dim % 32 == 0, "head_dim must be multiple of 32");
    // O matches Q's original layout
    auto O = torch::empty_strided(q.sizes(), q.strides(), q.options());
    auto O_view = (layout == 1) ? O.transpose(1, 2) : O;
    p.o = (bf16*)O_view.data_ptr();
    DISPATCH_HEAD_DIM(p.head_dim, dispatch_decode, p);
    return O;
 }
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("attn_decode", &attn_decode,
        py::arg("q"),
        py::arg("k"),
        py::arg("v"),
        py::arg("mask") = py::none(),
        py::arg("causal_offset") = -1,
        py::arg("scale") = 0.0,
        py::arg("layout") = 0,
        "GQA decode (tensor-core head-packing on sm_80+, scalar fallback)");
 }
@@ -0,0 +1,132 @@
 #pragma once
 #include <cuda_bf16.h>
 #include <float.h>
 #include "attn_common.h"
 using bf16 = __nv_bfloat16;
 constexpr int DC_CHUNK = 64;
 __device__ inline float warp_reduce_sum(float val) {
    for (int offset = 16; offset > 0; offset >>= 1)
        val += __shfl_xor_sync(0xFFFFFFFF, val, offset);
    return val;
 }
 __global__ void attn_decode_split_kv_kernel(AttentionParams<bf16> p) {
    int batch = blockIdx.x / p.kv_head;
    int kv_head = blockIdx.x % p.kv_head;
    int split = blockIdx.z;
    int group_size = blockDim.y;
    int q_head = kv_head * group_size + threadIdx.y;
    int lane = threadIdx.x;
    int hd_per_thread = p.head_dim / 32;
    // Q: [batch, q_head, q_len=1, head_dim] — stride-based
    float q_reg[8];
    int q_off = batch * p.q_stride_b + q_head * p.q_stride_h
              + lane * hd_per_thread * p.q_stride_d;
    for (int i = 0; i < hd_per_thread; i++)
        q_reg[i] = __bfloat162float(p.q[q_off + i * p.q_stride_d]);
    // KV: [batch, kv_head, kv_len, head_dim] — stride-based base
    int kv_base = batch * p.kv_stride_b + kv_head * p.kv_stride_h;
    int mask_base = batch * p.mask_b_stride;
    float m = -FLT_MAX, d = 0.0f, acc_reg[8] = {0.0f};
    extern __shared__ __align__(16) bf16 k_smem[];
    // Split-KV: each split processes a contiguous subset of chunks
    int chunks_total = (p.kv_len + DC_CHUNK - 1) / DC_CHUNK;
    int chunks_per_split = (chunks_total + p.num_splits - 1) / p.num_splits;
    int ch_begin = split * chunks_per_split;
    int ch_end = min(chunks_total, ch_begin + chunks_per_split);
    for (int ci = ch_begin; ci < ch_end; ci++) {
        int chunk_start = ci * DC_CHUNK;
        int this_chunk = min(DC_CHUNK, p.kv_len - chunk_start);
        // Load K into shared memory (gather from strided global)
        int total = this_chunk * p.head_dim;
        for (int i = threadIdx.y * 32 + lane; i < total; i += blockDim.x * blockDim.y) {
            int s = i / p.head_dim;
            int d_dim = i % p.head_dim;
            int kv_idx = chunk_start + s;
            int g_off = kv_base + kv_idx * p.kv_stride_l + d_dim * p.kv_stride_d;
            k_smem[i] = p.k[g_off];
        }
        __syncthreads();
        for (int s = 0; s < this_chunk; s++) {
            float partial = 0.0f;
            for (int i = 0; i < hd_per_thread; i++)
                partial += q_reg[i] * __bfloat162float(k_smem[s * p.head_dim + lane * hd_per_thread + i]);
            partial = warp_reduce_sum(partial) * p.scale;
            int kv_idx = chunk_start + s;
            if (p.use_mask && p.mask && !p.mask[mask_base + kv_idx])
                partial = -FLT_MAX;
            if (p.causal_offset >= 0 && kv_idx > p.causal_offset)
                partial = -FLT_MAX;
            float new_m = fmaxf(m, partial);
            float alpha = expf(m - new_m);
            float beta  = expf(partial - new_m);
            d = d * alpha + beta;
            // V: stride-based read
            int v_off = kv_base + kv_idx * p.kv_stride_l + lane * hd_per_thread * p.kv_stride_d;
            for (int i = 0; i < hd_per_thread; i++)
                acc_reg[i] = acc_reg[i] * alpha + __bfloat162float(p.v[v_off + i * p.kv_stride_d]) * beta;
            m = new_m;
        }
        __syncthreads();
    }
    // ---- write UN-normalised partials for this split ----
    size_t bh = (size_t)batch * p.q_head + q_head;
    size_t slot = bh * p.num_splits + split;
    int d0 = lane * hd_per_thread;
    for (int i = 0; i < hd_per_thread; i++) {
        int dd = d0 + i;
        p.o_part[slot * p.head_dim + dd] = acc_reg[i];
    }
    if (lane == 0) {
        p.ml_part[slot * 2] = m;
        p.ml_part[slot * 2 + 1] = d;
    }
 }
 // Reduce split-K partials into the final bf16 output. One block per (batch,
 // q_head); each thread folds across all splits with a single-pass
 // online-rescale reduction (expf + FMA counts halved vs 3-pass original).
 __global__ void attn_decode_combine_kernel(AttentionParams<bf16> p) {
    int bh = blockIdx.x;
    int d = threadIdx.x;
    if (d >= p.head_dim) return;
    int batch = bh / p.q_head;
    int q_head = bh % p.q_head;
    size_t split_base = (size_t)bh * p.num_splits;
    const float* mlp = p.ml_part + split_base * 2;
    const float* op = p.o_part + split_base * p.head_dim;
    float m = -FLT_MAX, l = 0.0f, acc = 0.0f;
    for (int s = 0; s < p.num_splits; s++) {
        float mi = mlp[s * 2];
        if (mi <= -FLT_MAX) continue;
        float li = mlp[s * 2 + 1];
        float nm = fmaxf(m, mi);
        float corr = __expf(m - nm);
        float e = __expf(mi - nm);
        acc = acc * corr + op[s * p.head_dim + d] * e;
        l = l * corr + li * e;
        m = nm;
    }
    float inv = (l > 1e-20f) ? (1.0f / l) : 0.0f;
    // Stride-based output write (q_len=1 for decode, so stride_l not needed)
    int o_off = batch * p.q_stride_b + q_head * p.q_stride_h + d * p.q_stride_d;
    p.o[o_off] = __float2bfloat16(acc * inv);
 }
@@ -0,0 +1,176 @@
 #pragma once
 #include <cfloat>
 #include <cuda_bf16.h>
 #include "attn_common.h"
 #include "attn_mma_utils.cuh"
 using bf16 = __nv_bfloat16;
 // Split-K (FlashDecoding) tensor-core decode via GQA head-packing.
 //
 // Decode has q_len == 1, so S = q @ K^T is a GEMV per head — no tensor-core
 // work on its own. But GQA gives us G = q_head / kv_head query heads that all
 // share one kv_head. We pack those G heads into the M=16 rows of
 // mma.sync.m16n8k16, turning G independent GEMVs into a single GEMM that
 // reuses each loaded K/V tile across all G heads (K/V load is the decode
 // bottleneck, so the reuse is the win, not the flops). The KV sequence is
 // partitioned across gridDim.z blocks so that a decode with only
 // batch*kv_head independent tasks can fill all SMs. Each (batch, kv_head,
 // split) block computes an UN-normalised partial (Oacc, m, l) over its KV
 // slice; the combine kernel below reduces across splits. Fixes the "grid too
 // small" bottleneck (0.04 waves/SM → many blocks) for long-context,
 // small-batch decode.
 template <int HEAD_DIM, int BC, int STAGES = 2>
 __global__ void attn_decode_split_kv_mma_kernel(AttentionParams<bf16> p) {
    constexpr int KD = HEAD_DIM / 16;
    constexpr int NC8 = BC / 8;
    constexpr int KT2 = BC / 16;
    constexpr int DN8 = HEAD_DIM / 8;
    constexpr int LD = HEAD_DIM;
    constexpr int SWIZ_MASK = (HEAD_DIM >= 64) ? 7 : (HEAD_DIM / 8 - 1);
    constexpr int VEC = 8;
    constexpr int TOTAL = BC * HEAD_DIM;
    const int lane = threadIdx.x;
    const int gid = lane >> 2;
    const int tid4 = lane & 3;
    const int kv_head = blockIdx.x;
    const int batch = blockIdx.y;
    const int split = blockIdx.z;
    const int G = p.q_head / p.kv_head;
    const int q_head0 = kv_head * G;
    // Double-buffered shared memory for K/V (no sQ needed — Q goes direct
    // from global to registers).
    __shared__ __align__(16) bf16 sK[STAGES * BC * LD];
    __shared__ __align__(16) bf16 sV[STAGES * BC * LD];
    // ---- Load Q directly from global into mma A-operand registers ----
    const int q_base = batch * p.q_stride_b + q_head0 * p.q_stride_h;
    const int qra = gid;
    const int qrb = gid + 8;
    const bool va = qra < G, vb = qrb < G;
    unsigned Qa[KD][4];
    load_q_mma_frags<KD>(p.q + q_base, p.q_stride_h, p.q_stride_d,
                          qra, qrb, va, vb, tid4, Qa);
    float Oacc[DN8][4];
 #pragma unroll
    for (int j = 0; j < DN8; j++)
        Oacc[j][0] = Oacc[j][1] = Oacc[j][2] = Oacc[j][3] = 0.0f;
    float m0 = -FLT_MAX, m1 = -FLT_MAX, l0 = 0.0f, l1 = 0.0f;
    // KV: stride-based base — [batch, kv_head, kv_len, head_dim]
    const int kv_base = batch * p.kv_stride_b + kv_head * p.kv_stride_h;
    const int tiles_total = (p.kv_len + BC - 1) / BC;
    const int tiles_per_split = (tiles_total + p.num_splits - 1) / p.num_splits;
    const int ti_begin = split * tiles_per_split;
    const int ti_end = min(tiles_total, ti_begin + tiles_per_split);
    const int has_mask = p.use_mask && p.mask;
    // ---- Load tile lambda: predicated cp.async, unified full/partial ----
    auto load_tile = [&](int ti, int buf) {
        int kv0 = ti * BC;
        bf16* dK = sK + buf * BC * LD;
        bf16* dV = sV + buf * BC * LD;
 #pragma unroll
        for (int i = lane * VEC; i < TOTAL; i += 32 * VEC) {
            int r = i / HEAD_DIM, d = i % HEAD_DIM;
            int kc = kv0 + r;
            bool valid = kc < p.kv_len;
            int off = r * LD + swiz_col(d, r, SWIZ_MASK);
            // KV stride-based: contiguous within head_dim (stride_d == 1 typically)
            int g_off = kv_base + kc * p.kv_stride_l + d * p.kv_stride_d;
            cp_async_16_pred(&dK[off], &p.k[g_off], valid);
            cp_async_16_pred(&dV[off], &p.v[g_off], valid);
        }
        cp_async_commit();
    };
    // ---- Prologue: issue first tile load ----
    if (ti_begin < ti_end) {
        load_tile(ti_begin, 0);
    }
    for (int ti = ti_begin; ti < ti_end; ti++) {
        constexpr int BUF_MASK = (STAGES > 1) ? (STAGES - 1) : 0;
        int buf = (ti - ti_begin) & BUF_MASK;
        // Wait for current tile, then issue next tile's prefetch (overlaps
        // with this tile's compute). Single syncwarp covers both hazards.
        // When STAGES==1, no prefetch — load happens at end of prior iter.
        cp_async_wait_group<0>();
        __syncwarp();
        if constexpr (STAGES > 1) {
            if (ti + 1 < ti_end)
                load_tile(ti + 1, (ti + 1 - ti_begin) & BUF_MASK);
        }
        const bf16* bK = sK + buf * BC * LD;
        const bf16* bV = sV + buf * BC * LD;
        int kv0 = ti * BC;
        float Sacc[NC8][4];
        mma_compute_scores<KD, NC8>(Qa, bK, LD, SWIZ_MASK, lane, Sacc);
        #pragma unroll
        for (int n8 = 0; n8 < NC8; n8++)
            Sacc[n8][0] *= p.scale, Sacc[n8][1] *= p.scale,
            Sacc[n8][2] *= p.scale, Sacc[n8][3] *= p.scale;
        // Decode: q_len=1, so qrow0=qrow1=0, mask_q_stride irrelevant
        int maxc = (p.causal_offset >= 0) ? min(p.kv_len, p.causal_offset + 1) : p.kv_len;
        mma_softmax_tile<NC8, DN8>(kv0, maxc, maxc,
                                    0, 0,
                                    p.mask_b_stride, 0,
                                    batch,
                                    p.mask, has_mask,
                                    Sacc, Oacc, m0, m1, l0, l1, lane);
        mma_pv_accumulate<DN8, KT2>(Sacc, bV, LD, SWIZ_MASK, lane, Oacc);
        __syncwarp();
        if constexpr (STAGES == 1) {
            if (ti + 1 < ti_end)
                load_tile(ti + 1, 0);
        }
    }
    // ---- write UN-normalised partials for this split ----
    auto split_slot = [&](int h) -> size_t {
        size_t bh = (size_t)batch * p.q_head + h;
        return bh * p.num_splits + split;
    };
 #pragma unroll
    for (int dn8 = 0; dn8 < DN8; dn8++) {
        int d = dn8 * 8 + 2 * tid4;
        int r0 = gid, r1 = gid + 8;
        if (r0 < G) {
            int h = q_head0 + r0;
            float* op = p.o_part + split_slot(h) * HEAD_DIM;
            op[d] = Oacc[dn8][0];
            op[d + 1] = Oacc[dn8][1];
        }
        if (r1 < G) {
            int h = q_head0 + r1;
            float* op = p.o_part + split_slot(h) * HEAD_DIM;
            op[d] = Oacc[dn8][2];
            op[d + 1] = Oacc[dn8][3];
        }
    }
    if (tid4 == 0) {
        int r0 = gid, r1 = gid + 8;
        if (r0 < G) {
            int h = q_head0 + r0;
            float* mp = p.ml_part + split_slot(h) * 2;
            mp[0] = m0; mp[1] = l0;
        }
        if (r1 < G) {
            int h = q_head0 + r1;
            float* mp = p.ml_part + split_slot(h) * 2;
            mp[0] = m1; mp[1] = l1;
        }
    }
 }
@@ -0,0 +1,177 @@
 #pragma once
 #include <torch/extension.h>
 #include <c10/cuda/CUDAGuard.h>
 #include "attn_common.h"
 using bf16 = __nv_bfloat16;
 inline int compute_num_splits(int base_blocks, int tiles_total) {
    int sm_count = 0;
    cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, 0);
    int n = (2 * sm_count + base_blocks - 1) / base_blocks;
    return std::max(1, std::min(n, std::min(tiles_total, 32)));
 }
 // Dispatch head_dim: shared macro — avoids C++20 lambda template syntax.
 // Usage: DISPATCH_HEAD_DIM(hd, fn, arg)
 //   Expands to: fn<32>(arg); fn<64>(arg); etc.
 #define DISPATCH_HEAD_DIM(hd, fn, arg) \
    switch (hd) { \
        case 32:  fn<32>(arg); break; \
        case 64:  fn<64>(arg); break; \
        case 128: fn<128>(arg); break; \
        case 256: fn<256>(arg); break; \
        default: \
            TORCH_CHECK(false, "unsupported head_dim ", hd, \
                         " (supported: 32, 64, 128, 256)"); \
    }
 template<typename P>
 inline void alloc_split_partials(P& p) {
    auto fopt = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA);
    auto o_part = torch::empty({p.batch, p.q_head, p.num_splits, p.head_dim}, fopt);
    auto ml_part = torch::empty({p.batch, p.q_head, p.num_splits, 2}, fopt);
    p.o_part = (float*)o_part.data_ptr();
    p.ml_part = (float*)ml_part.data_ptr();
 }
 // ---- Shared Q-dims + strides extraction ----
 template <typename P>
 inline void extract_q_dims_and_strides(torch::Tensor& q, int64_t layout, P& p) {
    if (layout == 1) q = q.transpose(1, 2);
    p.batch = (int)q.size(0);
    p.q_head = (int)q.size(1);
    p.q_len = (int)q.size(2);
    p.head_dim = (int)q.size(3);
    p.q_stride_b = (int)q.stride(0);
    p.q_stride_h = (int)q.stride(1);
    p.q_stride_l = (int)q.stride(2);
    p.q_stride_d = (int)q.stride(3);
 }
 // ---- Shared mask packing ----
 template <typename P>
 inline void pack_mask(const c10::optional<torch::Tensor>& mask, P& p) {
    if (p.use_mask) {
        auto m = mask.value();
        TORCH_CHECK(m.is_cuda(), "mask must be on CUDA");
        TORCH_CHECK(m.dtype() == torch::kBool, "mask must be bool");
        TORCH_CHECK(m.size(0) == p.batch, "mask batch mismatch");
        TORCH_CHECK(m.size(m.dim() - 1) == p.kv_len, "mask kv_len mismatch");
        if (m.dim() == 2) {
            p.mask_b_stride = (int)m.stride(0);
            p.mask_q_stride = 0;
        } else if (m.dim() == 3) {
            TORCH_CHECK(m.size(1) == p.q_len, "mask q_len mismatch");
            p.mask_b_stride = (int)m.stride(0);
            p.mask_q_stride = (int)m.stride(1);
        } else {
            TORCH_CHECK(false, "mask must be 2D [batch, kv_len] or 3D [batch, q_len, kv_len]");
        }
        p.mask = m.data_ptr<bool>();
    } else {
        p.mask = nullptr;
        p.mask_b_stride = 0;
        p.mask_q_stride = 0;
    }
 }
 // ---- attn_pack_params (contiguous KV) ----
 template<typename T>
 inline void attn_pack_params(
    torch::Tensor q,
    torch::Tensor k,
    torch::Tensor v,
    c10::optional<torch::Tensor> mask,
    int64_t causal_offset,
    double scale,
    int64_t layout,
    AttentionParams<T>& p
 ) {
    const at::cuda::OptionalCUDAGuard device_guard(device_of(q));
    TORCH_CHECK(q.is_cuda() && k.is_cuda() && v.is_cuda());
    TORCH_CHECK(q.dtype() == torch::kBFloat16);
    TORCH_CHECK(k.dtype() == torch::kBFloat16);
    TORCH_CHECK(v.dtype() == torch::kBFloat16);
    TORCH_CHECK(k.sizes() == v.sizes(), "K and V must have identical shapes");
    TORCH_CHECK(q.dim() == 4 && k.dim() == 4, "Q/K/V must be 4D");
    extract_q_dims_and_strides(q, layout, p);
    if (layout == 1) k = k.transpose(1, 2), v = v.transpose(1, 2);
    p.kv_head = (int)k.size(1);
    p.kv_len = (int)k.size(2);
    TORCH_CHECK(k.size(3) == p.head_dim, "K/V head_dim must match Q");
    p.kv_stride_b = (int)k.stride(0);
    p.kv_stride_h = (int)k.stride(1);
    p.kv_stride_l = (int)k.stride(2);
    p.kv_stride_d = (int)k.stride(3);
    p.causal_offset = (int)causal_offset;
    p.use_mask = mask.has_value() ? 1 : 0;
    p.scale = (scale > 0.0) ? (float)scale : 1.0f / sqrtf((float)p.head_dim);
    p.q = (const T*)q.data_ptr();
    p.k = (const T*)k.data_ptr();
    p.v = (const T*)v.data_ptr();
    p.o = nullptr;
    p.o_part = nullptr;
    p.ml_part = nullptr;
    pack_mask(mask, p);
 }
 // ---- attn_pack_paged_params ----
 template<typename T>
 inline void attn_pack_paged_params(
    torch::Tensor q,
    torch::Tensor page_table,
    torch::Tensor k_cache,
    torch::Tensor v_cache,
    int64_t page_size,
    int64_t kv_len,
    c10::optional<torch::Tensor> mask,
    int64_t causal_offset,
    double scale,
    int64_t layout,
    PagedAttentionParams<T>& p
 ) {
    const at::cuda::OptionalCUDAGuard device_guard(device_of(q));
    TORCH_CHECK(q.is_cuda() && page_table.is_cuda() && k_cache.is_cuda() && v_cache.is_cuda());
    TORCH_CHECK(q.dtype() == torch::kBFloat16, "q must be bf16");
    TORCH_CHECK(k_cache.dtype() == torch::kBFloat16, "k_cache must be bf16");
    TORCH_CHECK(v_cache.dtype() == torch::kBFloat16, "v_cache must be bf16");
    TORCH_CHECK(page_table.dtype() == torch::kLong, "page_table must be int64");
    TORCH_CHECK(k_cache.sizes() == v_cache.sizes(), "k_cache and v_cache must have identical shapes");
    extract_q_dims_and_strides(q, layout, p);
    p.kv_head = (int)k_cache.size(2);
    p.kv_len = (int)kv_len;
    p.page_size = (int)page_size;
    p.max_pages = (int)page_table.size(1);
    TORCH_CHECK(q.size(2) == 1, "Q seq_len must be 1 (decode)");
    TORCH_CHECK(p.head_dim % 32 == 0, "head_dim must be multiple of 32");
    TORCH_CHECK(k_cache.size(1) == page_size,
                "k_cache dim 1 must equal page_size, got ",
                k_cache.size(1), " vs ", page_size);
    p.causal_offset = (int)causal_offset;
    p.use_mask = (mask.has_value() && mask.value().defined()) ? 1 : 0;
    p.scale = (scale > 0.0) ? (float)scale : 1.0f / sqrtf((float)p.head_dim);
    p.page_table = page_table.data_ptr<int64_t>();
    p.k_cache = (const T*)k_cache.data_ptr();
    p.v_cache = (const T*)v_cache.data_ptr();
    p.q = (const T*)q.data_ptr();
    p.o = nullptr;
    p.o_part = nullptr;
    p.ml_part = nullptr;
    pack_mask(mask, p);
 }
@@ -0,0 +1,293 @@
 #pragma once
 #include <cfloat>
 #include <cuda_fp16.h>
 #include <cuda_runtime.h>
 // Shared MMA utilities for tensor-core GQA kernels.
 // mma.sync.m16n8k16 PTX wrappers, ldmatrix helpers, and bf16 packing.
 // mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32
 __device__ __forceinline__ void mma16816(float* d, const unsigned* a,
                                          const unsigned* b, const float* c) {
    asm volatile(
        "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
        "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};"
        : "=f"(d[0]), "=f"(d[1]), "=f"(d[2]), "=f"(d[3])
        : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]),
          "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));
 }
 // read two adjacent bf16 from smem as one packed .b32 (elem0 low, elem1 high)
 __device__ __forceinline__ unsigned ld2(const bf16* p) {
    return *reinterpret_cast<const unsigned*>(p);
 }
 // pack two floats into one bf16x2 as .b32
 __device__ __forceinline__ unsigned pk2(float a, float b) {
    __nv_bfloat162 v = __floats2bfloat162_rn(a, b);
    return *reinterpret_cast<unsigned*>(&v);
 }
 // pack two (non-contiguous) bf16 into one .b32
 __device__ __forceinline__ unsigned pkb(bf16 a, bf16 b) {
    __nv_bfloat162 v;
    v.x = a;
    v.y = b;
    return *reinterpret_cast<unsigned*>(&v);
 }
 // ldmatrix: cooperatively load mma fragments from smem (one instruction per
 // 16x16 / 16x8 tile) with the exact register layout mma expects — replaces the
 // scalar per-thread fragment packing, cutting shared-load instructions and bank
 // conflicts. Each lane supplies the shared address of one 8-wide row.
 __device__ __forceinline__ void ldmatrix_x4(unsigned* r, const bf16* p) {
    unsigned a = __cvta_generic_to_shared(p);
    asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];"
                 : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3])
                 : "r"(a));
 }
 __device__ __forceinline__ void ldmatrix_x2(unsigned* r, const bf16* p) {
    unsigned a = __cvta_generic_to_shared(p);
    asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0,%1}, [%2];"
                 : "=r"(r[0]), "=r"(r[1])
                 : "r"(a));
 }
 __device__ __forceinline__ void ldmatrix_x2_trans(unsigned* r, const bf16* p) {
    unsigned a = __cvta_generic_to_shared(p);
    asm volatile("ldmatrix.sync.aligned.m8n8.x2.trans.shared.b16 {%0,%1}, [%2];"
                 : "=r"(r[0]), "=r"(r[1])
                 : "r"(a));
 }
 // XOR swizzle for shared-memory column at 8-bf16 chunk granularity.
 // Eliminates ldmatrix bank conflicts without LD padding: consecutive rows
 // land in distinct bank groups. swiz_col(d, r, mask) = ((d>>3)^(r&mask))<<3 | (d&7).
 // mask must cover log2(HEAD_DIM/8) chunk bits but stay within LD: use 7 for
 // HEAD_DIM>=64 (8+ chunks), 3 for HEAD_DIM=32 (4 chunks). Default 7 keeps
 // existing HEAD_DIM>=64 call sites working unchanged.
 __device__ __forceinline__ int swiz_col(int d, int r, int mask = 7) {
    return ((d >> 3) ^ (r & mask)) << 3 | (d & 7);
 }
 // cp.async: copy 16 bytes (8 bf16) from global to shared memory directly,
 // bypassing registers. Eliminates shared-store bank conflicts and cuts
 // load-loop instruction count in half (1 cp.async vs 1 LDG + 1 STS).
 // Requires sm_80+.
 __device__ __forceinline__ void cp_async_16(bf16* smem_ptr, const void* gmem_ptr) {
    unsigned smem_addr = __cvta_generic_to_shared(smem_ptr);
    asm volatile("cp.async.ca.shared.global [%0], [%1], 16;"
                 :: "r"(smem_addr), "l"(gmem_ptr));
 }
 // Predicated cp.async: copy 16 bytes when `pred`, otherwise zero-fill the
 // destination (src-size operand = 0 → no bytes read from src, so an
 // out-of-bounds src address is never dereferenced). Lets full and partial
 // tiles share one uniform async load path — no scalar fallback branch.
 __device__ __forceinline__ void cp_async_16_pred(bf16* smem_ptr,
                                                  const void* gmem_ptr,
                                                  bool pred) {
    unsigned smem_addr = __cvta_generic_to_shared(smem_ptr);
    int src_size = pred ? 16 : 0;
    asm volatile("cp.async.ca.shared.global [%0], [%1], 16, %2;"
                 :: "r"(smem_addr), "l"(gmem_ptr), "r"(src_size));
 }
 __device__ __forceinline__ void cp_async_commit() {
    asm volatile("cp.async.commit_group;");
 }
 __device__ __forceinline__ void cp_async_wait_all() {
    asm volatile("cp.async.wait_all;");
 }
 // Wait until at most N commit groups are still in flight. Used for
 // double-buffered pipelining: wait_group<1> lets the next tile's cp.async
 // continue while ensuring the current tile's data is ready.
 template <int N>
 __device__ __forceinline__ void cp_async_wait_group() {
    asm volatile("cp.async.wait_group %0;" :: "n"(N));
 }
 // ---------------------------------------------------------------------------
 // Q-load: load query rows directly from global memory into mma A-operand
 // register layout. One call replaces ~15 duplicated lines in each MMA kernel.
 // stride_row is p.q_stride_h for decode (q_len=1, G heads) or
 //              p.q_stride_l for prefill (multi-q rows).
 // ---------------------------------------------------------------------------
 template <int KD>
 __device__ inline void load_q_mma_frags(
    const bf16* __restrict__ q,
    int stride_row,
    int stride_d,
    int qra, int qrb,
    bool va, bool vb,
    int tid4,
    unsigned Qa[KD][4])
 {
    #pragma unroll
    for (int kt = 0; kt < KD; kt++) {
        int c = kt * 16 + tid4 * 2;
        const unsigned* pau = reinterpret_cast<const unsigned*>(
            &q[qra * stride_row + c * stride_d]);
        const unsigned* pbu = reinterpret_cast<const unsigned*>(
            &q[qrb * stride_row + c * stride_d]);
        Qa[kt][0] = va ? pau[0] : 0u;
        Qa[kt][1] = vb ? pbu[0] : 0u;
        Qa[kt][2] = va ? pau[4] : 0u;
        Qa[kt][3] = vb ? pbu[4] : 0u;
    }
 }
 // ---------------------------------------------------------------------------
 // Shared MMA compute functions — used by both decode and prefill MMA kernels.
 // Extracted because S=Q@K^T, online softmax, and P@V are structurally identical
 // between the two kernels; only the per-row causal/mask bounds differ.
 // ---------------------------------------------------------------------------
 // S = Q @ K^T  (Qa pre-loaded by the caller; scale applied post-mma in the
 // caller to avoid bf16 precision loss).
 // LD and SWIZ_MASK are constexpr in the calling kernel — passing them as
 // runtime ints lets the compiler fold them while keeping the signature clean.
 template <int KD, int NC8>
 __device__ inline void mma_compute_scores(
    const unsigned Qa[KD][4],
    const bf16* __restrict__ sK,
    int LD, 
    int SWIZ_MASK, 
    int lane,
    float Sacc[NC8][4])
 {
    #pragma unroll
    for (int n8 = 0; n8 < NC8; n8++) {
        Sacc[n8][0] = Sacc[n8][1] = Sacc[n8][2] = Sacc[n8][3] = 0.0f;
        int krow_l = n8 * 8 + (lane & 7);
        int kcol_h = (lane & 8) ? 8 : 0;
        #pragma unroll
        for (int kt = 0; kt < KD; kt++) {
            unsigned b[2];
            ldmatrix_x2(b, &sK[krow_l * LD + swiz_col(kt * 16 + kcol_h, krow_l, SWIZ_MASK)]);
            mma16816(Sacc[n8], Qa[kt], b, Sacc[n8]);
        }
    }
 }
 // Online softmax + Oacc rescale for one K/V tile.
 //   maxc0/maxc1: per-row KV column bounds (prefill: per-query-row causal limits;
 //     decode: same value for both rows since q_len==1).
 //   qrow0/qrow1: query row indices (for 3D mask indexing; decode passes 0).
 //   mask_b_stride/mask_q_stride: mask layout (2D: mask_q_stride=0; 3D: =kv_len).
 // Reads Sacc (Q@K^T scores), applies causal/mask, computes P = exp(S - nm),
 // rescales Oacc by exp(m_old - nm), and updates m/l — all in place.
 template <int NC8, int DN8>
 __device__ inline void mma_softmax_tile(
    int kv0,
    int maxc0, 
    int maxc1,
    int qrow0, 
    int qrow1,
    int mask_b_stride, 
    int mask_q_stride,
    int mask_batch,
    const bool* __restrict__ mask,
    bool has_mask,
    float Sacc[NC8][4],
    float Oacc[DN8][4],
    float& m0, float& m1,
    float& l0, float& l1,
    int lane)
 {
    int tid4 = lane & 3;
    // Mask out-of-bounds / masked columns: set -FLT_MAX so expf → 0 downstream
    // without per-element sentinel checks. Compute tile-local row maxima.
    float rmax0 = -FLT_MAX, rmax1 = -FLT_MAX;
    int mask_base0 = mask_batch * mask_b_stride + qrow0 * mask_q_stride;
    int mask_base1 = mask_batch * mask_b_stride + qrow1 * mask_q_stride;
    #pragma unroll
    for (int n8 = 0; n8 < NC8; n8++) {
        int cc = kv0 + n8 * 8 + 2 * tid4;
        int c1 = cc + 1;
        bool b0 = (cc >= maxc0) || (has_mask && !mask[mask_base0 + cc]);
        bool b1 = (c1 >= maxc0) || (has_mask && !mask[mask_base0 + c1]);
        bool b2 = (cc >= maxc1) || (has_mask && !mask[mask_base1 + cc]);
        bool b3 = (c1 >= maxc1) || (has_mask && !mask[mask_base1 + c1]);
        float s0 = b0 ? -FLT_MAX : Sacc[n8][0];
        float s1 = b1 ? -FLT_MAX : Sacc[n8][1];
        float s2 = b2 ? -FLT_MAX : Sacc[n8][2];
        float s3 = b3 ? -FLT_MAX : Sacc[n8][3];
        Sacc[n8][0] = s0; Sacc[n8][1] = s1;
        Sacc[n8][2] = s2; Sacc[n8][3] = s3;
        rmax0 = fmaxf(rmax0, fmaxf(s0, s1));
        rmax1 = fmaxf(rmax1, fmaxf(s2, s3));
    }
    // Warp-reduce row maxima across the 4-lane thread group (xor 1, xor 2).
    rmax0 = fmaxf(rmax0, __shfl_xor_sync(0xFFFFFFFF, rmax0, 1));
    rmax0 = fmaxf(rmax0, __shfl_xor_sync(0xFFFFFFFF, rmax0, 2));
    rmax1 = fmaxf(rmax1, __shfl_xor_sync(0xFFFFFFFF, rmax1, 1));
    rmax1 = fmaxf(rmax1, __shfl_xor_sync(0xFFFFFFFF, rmax1, 2));
    // nm = max(running max m, tile-local max rmax) — updated running maximum.
    float nm0 = fmaxf(m0, rmax0), nm1 = fmaxf(m1, rmax1);
    // corr rescales Oacc and l by exp(m_old - nm). When all-masked (m == nm ==
    // -FLT_MAX), exp(0) = 1 — correct, no guard needed.
    float corr0 = __expf(m0 - nm0);
    float corr1 = __expf(m1 - nm1);
    // pn guards only the all-masked-row edge: if nm == -FLT_MAX, exp(S - nm)
    // gives 1 not 0 for masked entries. Two scalar masks replace 4*NC8
    // per-element comparisons.
    float pn0 = (nm0 == -FLT_MAX) ? 0.0f : 1.0f;
    float pn1 = (nm1 == -FLT_MAX) ? 0.0f : 1.0f;
    // P = exp(S - nm) for each element. Masked entries (Sacc = -FLT_MAX) give
    // exp(-inf) ≈ 0 naturally; pn zero-fills the all-masked-row edge.
    float rsum0 = 0.0f, rsum1 = 0.0f;
    #pragma unroll
    for (int n8 = 0; n8 < NC8; n8++) {
        float p0 = pn0 * __expf(Sacc[n8][0] - nm0);
        float p1 = pn0 * __expf(Sacc[n8][1] - nm0);
        float p2 = pn1 * __expf(Sacc[n8][2] - nm1);
        float p3 = pn1 * __expf(Sacc[n8][3] - nm1);
        Sacc[n8][0] = p0; Sacc[n8][1] = p1;
        Sacc[n8][2] = p2; Sacc[n8][3] = p3;
        rsum0 += p0 + p1;
        rsum1 += p2 + p3;
    }
    rsum0 += __shfl_xor_sync(0xFFFFFFFF, rsum0, 1);
    rsum0 += __shfl_xor_sync(0xFFFFFFFF, rsum0, 2);
    rsum1 += __shfl_xor_sync(0xFFFFFFFF, rsum1, 1);
    rsum1 += __shfl_xor_sync(0xFFFFFFFF, rsum1, 2);
    l0 = l0 * corr0 + rsum0;
    l1 = l1 * corr1 + rsum1;
    m0 = nm0; m1 = nm1;
    #pragma unroll
    for (int j = 0; j < DN8; j++) {
        Oacc[j][0] *= corr0; Oacc[j][1] *= corr0;
        Oacc[j][2] *= corr1; Oacc[j][3] *= corr1;
    }
 }
 // O += P @ V  (Sacc must contain P = attention weights after softmax).
 template <int DN8, int KT2>
 __device__ inline void mma_pv_accumulate(
    float Sacc[][4],
    const bf16* __restrict__ sV,
    int LD, int SWIZ_MASK, int lane,
    float Oacc[DN8][4])
 {
    #pragma unroll
    for (int kt2 = 0; kt2 < KT2; kt2++) {
        unsigned Pa[4];
        Pa[0] = pk2(Sacc[kt2 * 2][0], Sacc[kt2 * 2][1]);
        Pa[1] = pk2(Sacc[kt2 * 2][2], Sacc[kt2 * 2][3]);
        Pa[2] = pk2(Sacc[kt2 * 2 + 1][0], Sacc[kt2 * 2 + 1][1]);
        Pa[3] = pk2(Sacc[kt2 * 2 + 1][2], Sacc[kt2 * 2 + 1][3]);
        int vrow_l = kt2 * 16 + (lane & 15);
        #pragma unroll
        for (int dn8 = 0; dn8 < DN8; dn8++) {
            unsigned b[2];
            ldmatrix_x2_trans(b, &sV[vrow_l * LD + swiz_col(dn8 * 8, vrow_l, SWIZ_MASK)]);
            mma16816(Oacc[dn8], Pa, b, Oacc[dn8]);
        }
    }
 }
@@ -0,0 +1,82 @@
 #include "attn_paged_decode_split_kv.cuh"
 #ifndef ASTRAI_NO_MMA
 #include "attn_paged_decode_split_kv_mma.cuh"
 #endif
 #include "attn_entry_utils.cuh"
 static void launch_paged_scalar_decode(PagedAttentionParams<bf16>& p) {
    int group_size = p.q_head / p.kv_head;
    int chunks_total = (p.kv_len + PDC_CHUNK - 1) / PDC_CHUNK;
    p.num_splits = compute_num_splits(p.batch * p.kv_head, chunks_total);
    alloc_split_partials(p);
    size_t smem = PDC_CHUNK * p.head_dim * sizeof(bf16);
    dim3 grid = dim3(p.batch * p.kv_head, 1, p.num_splits);
    dim3 block = dim3(32, group_size);
    paged_attn_decode_split_kv_kernel<<<grid, block, smem>>>(p);
    paged_attn_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim>>>(p);
 }
 #ifndef ASTRAI_NO_MMA
 template <int HEAD_DIM, int BC, int STAGES = (HEAD_DIM <= 128) ? 2 : 1>
 static void launch_paged_mma_decode(PagedAttentionParams<bf16>& p) {
    int tiles_total = (p.kv_len + BC - 1) / BC;
    p.num_splits = compute_num_splits(p.batch * p.kv_head, tiles_total);
    alloc_split_partials(p);
    paged_attn_decode_split_kv_mma_kernel<HEAD_DIM, BC, STAGES><<<dim3(p.kv_head, p.batch, p.num_splits), 32>>>(p);
    paged_attn_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim>>>(p);
 }
 #endif
 template <int HEAD_DIM>
 static void dispatch_paged_decode(PagedAttentionParams<bf16>& p) {
 #ifndef ASTRAI_NO_MMA
    int G = p.q_head / p.kv_head;
    if (G >= 1 && G <= 16 && p.page_size >= 32) {
        launch_paged_mma_decode<HEAD_DIM, 32>(p);
        return;
    }
 #endif
    launch_paged_scalar_decode(p);
 }
 torch::Tensor attn_paged_decode(
    torch::Tensor q,
    torch::Tensor page_table,
    torch::Tensor k_cache,
    torch::Tensor v_cache,
    int64_t page_size,
    int64_t kv_len,
    c10::optional<torch::Tensor> mask,
    int64_t causal_offset,
    double scale,
    int64_t layout
 ) {
    PagedAttentionParams<bf16> p;
    attn_pack_paged_params(q, page_table, k_cache, v_cache,
                           page_size, kv_len, mask, causal_offset, scale, layout, p);
    auto O = torch::empty_strided(q.sizes(), q.strides(), q.options());
    auto O_view = (layout == 1) ? O.transpose(1, 2) : O;
    p.o = (bf16*)O_view.data_ptr();
    DISPATCH_HEAD_DIM(p.head_dim, dispatch_paged_decode, p);
    return O;
 }
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("attn_paged_decode", &attn_paged_decode,
        py::arg("q"),
        py::arg("page_table"),
        py::arg("k_cache"),
        py::arg("v_cache"),
        py::arg("page_size"),
        py::arg("kv_len"),
        py::arg("mask") = py::none(),
        py::arg("causal_offset") = -1,
        py::arg("scale") = 0.0,
        py::arg("layout") = 0,
        "Paged GQA decode — split-KV with direct page-table access.");
 }
@@ -0,0 +1,147 @@
 #pragma once
 #include <cuda_bf16.h>
 #include <float.h>
 #include "attn_common.h"
 using bf16 = __nv_bfloat16;
 constexpr int PDC_CHUNK = 64;
 __device__ inline float paged_warp_reduce_sum(float val) {
    for (int offset = 16; offset > 0; offset >>= 1)
        val += __shfl_xor_sync(0xFFFFFFFF, val, offset);
    return val;
 }
 // Split-KV scalar decode: one warp per query head, grid.z partitions KV.
 __global__ void paged_attn_decode_split_kv_kernel(PagedAttentionParams<bf16> p) {
    int batch = blockIdx.x / p.kv_head;
    int kv_head = blockIdx.x % p.kv_head;
    int split = blockIdx.z;
    int group_size = blockDim.y;
    int q_head = kv_head * group_size + threadIdx.y;
    int lane = threadIdx.x;
    int hd_per_thread = p.head_dim / 32;
    // Q: stride-based [batch, q_head, q_len=1, head_dim]
    float q_reg[8];
    int q_off = batch * p.q_stride_b + q_head * p.q_stride_h
              + lane * hd_per_thread * p.q_stride_d;
    #pragma unroll
    for (int i = 0; i < hd_per_thread; i++)
        q_reg[i] = __bfloat162float(p.q[q_off + i * p.q_stride_d]);
    float m = -FLT_MAX, d = 0.0f, acc_reg[8] = {0.0f};
    extern __shared__ __align__(16) bf16 k_smem[];
    int chunks_total = (p.kv_len + PDC_CHUNK - 1) / PDC_CHUNK;
    int chunks_per_split = (chunks_total + p.num_splits - 1) / p.num_splits;
    int ch_begin = split * chunks_per_split;
    int ch_end = min(chunks_total, ch_begin + chunks_per_split);
    const int mask_base = batch * p.mask_b_stride;
    for (int ci = ch_begin; ci < ch_end; ci++) {
        int chunk_start = ci * PDC_CHUNK;
        int this_chunk = min(PDC_CHUNK, p.kv_len - chunk_start);
        int total = this_chunk * p.head_dim;
        for (int i = threadIdx.y * 32 + lane; i < total; i += blockDim.x * blockDim.y) {
            int s = i / p.head_dim;
            int d_dim = i % p.head_dim;
            int pos = chunk_start + s;
            int logical_page = pos / p.page_size;
            int page_offset = pos % p.page_size;
            int phys_page = p.page_table[batch * p.max_pages + logical_page];
            if (phys_page >= 0) {
                int64_t off = (int64_t)phys_page * p.page_size * p.kv_head * p.head_dim
                            + (int64_t)page_offset * p.kv_head * p.head_dim
                            + (int64_t)kv_head * p.head_dim
                            + d_dim;
                k_smem[i] = p.k_cache[off];
            } else {
                k_smem[i] = __float2bfloat16(0.0f);
            }
        }
        __syncthreads();
        for (int s = 0; s < this_chunk; s++) {
            float partial = 0.0f;
            #pragma unroll
            for (int i = 0; i < hd_per_thread; i++)
                partial += q_reg[i] * __bfloat162float(k_smem[s * p.head_dim + lane * hd_per_thread + i]);
            partial = paged_warp_reduce_sum(partial) * p.scale;
            int kv_idx = chunk_start + s;
            if (p.use_mask && p.mask && !p.mask[mask_base + kv_idx])
                partial = -FLT_MAX;
            if (p.causal_offset >= 0 && kv_idx > p.causal_offset)
                partial = -FLT_MAX;
            float new_m = fmaxf(m, partial);
            float alpha = expf(m - new_m);
            float beta  = expf(partial - new_m);
            d = d * alpha + beta;
            int pos = chunk_start + s;
            int logical_page = pos / p.page_size;
            int page_offset = pos % p.page_size;
            int phys_page = p.page_table[batch * p.max_pages + logical_page];
            if (phys_page >= 0) {
                int64_t v_base = (int64_t)phys_page * p.page_size * p.kv_head * p.head_dim
                               + (int64_t)page_offset * p.kv_head * p.head_dim
                               + (int64_t)kv_head * p.head_dim;
                #pragma unroll
                for (int i = 0; i < hd_per_thread; i++)
                    acc_reg[i] = acc_reg[i] * alpha + __bfloat162float(p.v_cache[v_base + lane * hd_per_thread + i]) * beta;
            } else {
                #pragma unroll
                for (int i = 0; i < hd_per_thread; i++)
                    acc_reg[i] = acc_reg[i] * alpha + 0.0f * beta;
            }
            m = new_m;
        }
        __syncthreads();
    }
    size_t bh = (size_t)batch * p.q_head + q_head;
    size_t slot = bh * p.num_splits + split;
    int d0 = lane * hd_per_thread;
    #pragma unroll
    for (int i = 0; i < hd_per_thread; i++)
        p.o_part[slot * p.head_dim + (d0 + i)] = acc_reg[i];
    if (lane == 0) {
        p.ml_part[slot * 2] = m;
        p.ml_part[slot * 2 + 1] = d;
    }
 }
 __global__ void paged_attn_decode_combine_kernel(PagedAttentionParams<bf16> p) {
    int bh = blockIdx.x;
    int d = threadIdx.x;
    if (d >= p.head_dim) return;
    int batch = bh / p.q_head;
    int q_head = bh % p.q_head;
    size_t split_base = (size_t)bh * p.num_splits;
    const float* mlp = p.ml_part + split_base * 2;
    const float* op = p.o_part + split_base * p.head_dim;
    float m = -FLT_MAX, l = 0.0f, acc = 0.0f;
    for (int s = 0; s < p.num_splits; s++) {
        float mi = mlp[s * 2];
        if (mi <= -FLT_MAX) continue;
        float li = mlp[s * 2 + 1];
        float nm = fmaxf(m, mi);
        float corr = __expf(m - nm);
        float e = __expf(mi - nm);
        acc = acc * corr + op[s * p.head_dim + d] * e;
        l = l * corr + li * e;
        m = nm;
    }
    float inv = (l > 1e-20f) ? (1.0f / l) : 0.0f;
    int o_off = batch * p.q_stride_b + q_head * p.q_stride_h + d * p.q_stride_d;
    p.o[o_off] = __float2bfloat16(acc * inv);
 }
@@ -0,0 +1,170 @@
 #pragma once
 #include <cfloat>
 #include <cuda_bf16.h>
 #include "attn_common.h"
 #include "attn_mma_utils.cuh"
 using bf16 = __nv_bfloat16;
 // Paged split-KV tensor-core decode via GQA head-packing.
 // Identical algorithm to attn_decode_split_kv_mma_kernel but reads K/V
 // directly from the page pool through a page table, eliminating the gather
 // copy.  Each tile (BC=32) fits within a single page (page_size >= 32), so
 // the page-table lookup happens once per tile for cp.async.
 template <int HEAD_DIM, int BC, int STAGES = (HEAD_DIM <= 128) ? 2 : 1>
 __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams<bf16> p) {
    constexpr int KD = HEAD_DIM / 16;
    constexpr int NC8 = BC / 8;
    constexpr int KT2 = BC / 16;
    constexpr int DN8 = HEAD_DIM / 8;
    constexpr int LD = HEAD_DIM;
    constexpr int SWIZ_MASK = (HEAD_DIM >= 64) ? 7 : (HEAD_DIM / 8 - 1);
    constexpr int VEC = 8;
    constexpr int TOTAL = BC * HEAD_DIM;
    const int lane = threadIdx.x;
    const int gid = lane >> 2;
    const int tid4 = lane & 3;
    const int kv_head = blockIdx.x;
    const int batch = blockIdx.y;
    const int split = blockIdx.z;
    const int G = p.q_head / p.kv_head;
    const int q_head0 = kv_head * G;
    __shared__ __align__(16) bf16 sK[STAGES * BC * LD];
    __shared__ __align__(16) bf16 sV[STAGES * BC * LD];
    // ---- Load Q directly from global into mma A-operand registers ----
    const int q_base = batch * p.q_stride_b + q_head0 * p.q_stride_h;
    const int qra = gid;
    const int qrb = gid + 8;
    const bool va = qra < G, vb = qrb < G;
    unsigned Qa[KD][4];
    load_q_mma_frags<KD>(p.q + q_base, p.q_stride_h, p.q_stride_d,
                          qra, qrb, va, vb, tid4, Qa);
    float Oacc[DN8][4];
 #pragma unroll
    for (int j = 0; j < DN8; j++)
        Oacc[j][0] = Oacc[j][1] = Oacc[j][2] = Oacc[j][3] = 0.0f;
    float m0 = -FLT_MAX, m1 = -FLT_MAX, l0 = 0.0f, l1 = 0.0f;
    const int tiles_total = (p.kv_len + BC - 1) / BC;
    const int tiles_per_split = (tiles_total + p.num_splits - 1) / p.num_splits;
    const int ti_begin = split * tiles_per_split;
    const int ti_end = min(tiles_total, ti_begin + tiles_per_split);
    const int has_mask = p.use_mask && p.mask;
    // Paged strides (constant for the block)
    const int64_t page_stride = (int64_t)p.page_size * p.kv_head * HEAD_DIM;
    const int64_t pos_stride  = (int64_t)p.kv_head * HEAD_DIM;
    const int64_t head_off    = (int64_t)kv_head * HEAD_DIM;
    // ---- Load tile lambda: predicated cp.async, paged addressing ----
    auto load_tile = [&](int ti, int buf) {
        int kv0 = ti * BC;
        bf16* dK = sK + buf * BC * LD;
        bf16* dV = sV + buf * BC * LD;
        int logical_page = kv0 / p.page_size;
        int phys_page = p.page_table[batch * p.max_pages + logical_page];
        bool page_valid = (phys_page >= 0);
 #pragma unroll
        for (int i = lane * VEC; i < TOTAL; i += 32 * VEC) {
            int r = i / HEAD_DIM, d = i % HEAD_DIM;
            int kc = kv0 + r;
            bool valid = (kc < p.kv_len) && page_valid;
            int page_off = kc % p.page_size;
            int64_t gmem_base = (int64_t)phys_page * page_stride
                              + (int64_t)page_off * pos_stride
                              + head_off;
            int off = r * LD + swiz_col(d, r, SWIZ_MASK);
            cp_async_16_pred(&dK[off], &p.k_cache[gmem_base + d], valid);
            cp_async_16_pred(&dV[off], &p.v_cache[gmem_base + d], valid);
        }
        cp_async_commit();
    };
    // ---- Prologue: issue first tile load ----
    if (ti_begin < ti_end) {
        load_tile(ti_begin, 0);
    }
    for (int ti = ti_begin; ti < ti_end; ti++) {
        constexpr int BUF_MASK = (STAGES > 1) ? (STAGES - 1) : 0;
        int buf = (ti - ti_begin) & BUF_MASK;
        cp_async_wait_group<0>();
        __syncwarp();
        if constexpr (STAGES > 1) {
            if (ti + 1 < ti_end)
                load_tile(ti + 1, (ti + 1 - ti_begin) & BUF_MASK);
        }
        const bf16* bK = sK + buf * BC * LD;
        const bf16* bV = sV + buf * BC * LD;
        int kv0 = ti * BC;
        float Sacc[NC8][4];
        mma_compute_scores<KD, NC8>(Qa, bK, LD, SWIZ_MASK, lane, Sacc);
 #pragma unroll
        for (int n8 = 0; n8 < NC8; n8++)
            Sacc[n8][0] *= p.scale, Sacc[n8][1] *= p.scale,
            Sacc[n8][2] *= p.scale, Sacc[n8][3] *= p.scale;
        // Decode: q_len=1, so qrow0=qrow1=0, mask_q_stride irrelevant
        int maxc = (p.causal_offset >= 0) ? min(p.kv_len, p.causal_offset + 1) : p.kv_len;
        mma_softmax_tile<NC8, DN8>(kv0, maxc, maxc,
                                    0, 0,
                                    p.mask_b_stride, 0,
                                    batch,
                                    p.mask, has_mask,
                                    Sacc, Oacc, m0, m1, l0, l1, lane);
        mma_pv_accumulate<DN8, KT2>(Sacc, bV, LD, SWIZ_MASK, lane, Oacc);
        __syncwarp();
        if constexpr (STAGES == 1) {
            if (ti + 1 < ti_end)
                load_tile(ti + 1, 0);
        }
    }
    // ---- write UN-normalised partials for this split ----
    auto split_slot = [&](int h) -> size_t {
        size_t bh = (size_t)batch * p.q_head + h;
        return bh * p.num_splits + split;
    };
 #pragma unroll
    for (int dn8 = 0; dn8 < DN8; dn8++) {
        int d = dn8 * 8 + 2 * tid4;
        int r0 = gid, r1 = gid + 8;
        if (r0 < G) {
            int h = q_head0 + r0;
            float* op = p.o_part + split_slot(h) * HEAD_DIM;
            op[d] = Oacc[dn8][0];
            op[d + 1] = Oacc[dn8][1];
        }
        if (r1 < G) {
            int h = q_head0 + r1;
            float* op = p.o_part + split_slot(h) * HEAD_DIM;
            op[d] = Oacc[dn8][2];
            op[d + 1] = Oacc[dn8][3];
        }
    }
    if (tid4 == 0) {
        int r0 = gid, r1 = gid + 8;
        if (r0 < G) {
            int h = q_head0 + r0;
            float* mp = p.ml_part + split_slot(h) * 2;
            mp[0] = m0; mp[1] = l0;
        }
        if (r1 < G) {
            int h = q_head0 + r1;
            float* mp = p.ml_part + split_slot(h) * 2;
            mp[0] = m1; mp[1] = l1;
        }
    }
 }
@@ -0,0 +1,64 @@
 #include "attn_prefill_split_q.cuh"
 #include "attn_entry_utils.cuh"
 #ifndef ASTRAI_NO_MMA
 #include "attn_prefill_split_q_mma.cuh"
 #endif
 template <int HEAD_DIM>
 static void dispatch_prefill(AttentionParams<bf16>& p) {
 #ifndef ASTRAI_NO_MMA
    constexpr int WARPS = 4, BR = 16;
    // KV tile: bigger tiles amortize the per-tile cp.async wait + barrier +
    // loop overhead over more tensor-core work (this kernel is latency-bound,
    // not compute/bandwidth-bound), so BC=32 wins ~6-8% over BC=16 for
    // D<=128. D=256 stays at 16: BC=32 double-buffered would need 64KB smem,
    // over the 48KB static cap.
    constexpr int BC = (HEAD_DIM <= 128) ? 32 : 16;
    dim3 grid((p.q_len + BR * WARPS - 1) / (BR * WARPS), p.q_head, p.batch);
    dim3 block(WARPS * 32, 1, 1);
    // Static shared memory — double-buffered K/V only (no sQ: Q goes direct
    // to registers). 2*BC*LD bf16 each for sK and sV → 4*BC*HEAD_DIM*2 bytes.
    // Occupancy is smem-capped: D=64→3 blocks/SM (16KB), D=128→1 (32KB),
    // D=256→1 (32KB, BC=16).
    attn_prefill_split_q_mma_kernel<HEAD_DIM, WARPS, BC><<<grid, block>>>(p);
 #else
    constexpr int G = 8, ROWS = 32, P_BC = 32;
    dim3 grid((p.q_len + ROWS - 1) / ROWS, p.q_head, p.batch);
    dim3 block(G, ROWS, 1);
    attn_prefill_split_q_kernel_t<HEAD_DIM, G, ROWS, P_BC><<<grid, block>>>(p);
 #endif
 }
 torch::Tensor attn_prefill(
    torch::Tensor q,
    torch::Tensor k,
    torch::Tensor v,
    c10::optional<torch::Tensor> mask,
    int64_t causal_offset,
    double scale,
    int64_t layout
 ) {
    AttentionParams<bf16> p;
    attn_pack_params(q, k, v, mask, causal_offset, scale, layout, p);
    TORCH_CHECK(p.head_dim % 16 == 0, "head_dim must be multiple of 16");
    auto O = torch::empty_strided(q.sizes(), q.strides(), q.options());
    auto O_view = (layout == 1) ? O.transpose(1, 2) : O;
    p.o = (bf16*)O_view.data_ptr();
    DISPATCH_HEAD_DIM(p.head_dim, dispatch_prefill, p);
    return O;
 }
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("attn_prefill", &attn_prefill,
        py::arg("q"),
        py::arg("k"),
        py::arg("v"),
        py::arg("mask") = py::none(),
        py::arg("causal_offset") = -1,
        py::arg("scale") = 0.0,
        py::arg("layout") = 0,
        "GQA prefill (tensor-core mma on sm_80+, scalar fallback)");
 }
@@ -0,0 +1,152 @@
 #pragma once
 #include <cfloat>
 #include <cuda_bf16.h>
 #include "attn_common.h"
 using bf16 = __nv_bfloat16;
 // v9: group-split register blocking. G threads cooperate on one query row,
 // each owning HEAD_DIM/G dims of qreg[]/acc[]. Small per-thread footprint keeps
 // occupancy high; the S dot product is reduced across the G-lane group with a
 // short shuffle chain (log2(G) shuffles) instead of a full 32-lane warp reduce.
 // Online (per-kv) softmax — cheap because acc[] is only HEAD_DIM/G long.
 // Templated on <HEAD_DIM, G, ROWS, P_BC>. Block = (G, ROWS). G power-of-two,
 // G*ROWS a multiple of 32 with groups warp-aligned.
 template <int G>
 __device__ __forceinline__ float group_reduce_sum(float v, unsigned mask) {
 #pragma unroll
    for (int o = G / 2; o > 0; o >>= 1)
        v += __shfl_xor_sync(mask, v, o);
    return v;
 }
 // load 8 contiguous bf16 from (16-byte aligned) smem as one float4, unpack to
 // 8 floats — cuts shared-load instructions 8x vs scalar bf16 loads.
 __device__ __forceinline__ void ld8(const bf16* p, float* o) {
    float4 raw = *reinterpret_cast<const float4*>(p);
    const __nv_bfloat162* h = reinterpret_cast<const __nv_bfloat162*>(&raw);
 #pragma unroll
    for (int j = 0; j < 4; j++) {
        float2 f = __bfloat1622float2(h[j]);
        o[2 * j] = f.x;
        o[2 * j + 1] = f.y;
    }
 }
 template <int HEAD_DIM, int G, int ROWS, int P_BC>
 __global__ void attn_prefill_split_q_kernel_t(AttentionParams<bf16> p) {
    constexpr int DPT = HEAD_DIM / G;
    int q_tile = blockIdx.x;
    int q_head = blockIdx.y;
    int batch  = blockIdx.z;
    int gpos   = threadIdx.x;  // 0..G-1  (which d-chunk)
    int row    = threadIdx.y;  // 0..ROWS-1
    int q_row  = q_tile * ROWS + row;
    int kv_head = q_head / (p.q_head / p.kv_head);
    __shared__ __align__(16) bf16 sK[P_BC * HEAD_DIM];
    __shared__ __align__(16) bf16 sV[P_BC * HEAD_DIM];
    // Q: stride-based load [batch, q_head, q_len, head_dim]
    float qreg[DPT];
    if (q_row < p.q_len) {
        int q_off = batch * p.q_stride_b + q_head * p.q_stride_h
                  + q_row * p.q_stride_l + gpos * DPT * p.q_stride_d;
 #pragma unroll
        for (int i = 0; i < DPT; i++)
            qreg[i] = __bfloat162float(p.q[q_off + i * p.q_stride_d]) * p.scale;
    }
    float m = -FLT_MAX, l = 0.0f;
    float acc[DPT];
 #pragma unroll
    for (int i = 0; i < DPT; i++)
        acc[i] = 0.0f;
    // KV: stride-based base
    int kv_base = batch * p.kv_stride_b + kv_head * p.kv_stride_h;
    int mask_batch_base = batch * p.mask_b_stride;
    int tiles   = (p.kv_len + P_BC - 1) / P_BC;
    int tt      = G * ROWS;
    int lid     = row * G + gpos;
    // per-group shuffle mask: only the G lanes of this row's group participate,
    // so causal masking (differing loop bounds across rows in a warp) is safe.
    int lane_in_warp = lid & 31;
    unsigned gmask = (G == 32) ? 0xFFFFFFFFu
                               : (((1u << G) - 1u) << (lane_in_warp & ~(G - 1)));
    for (int ti = 0; ti < tiles; ti++) {
        int kv0  = ti * P_BC;
        int tlen = min(P_BC, p.kv_len - kv0);
        // Load K/V into shared memory from strided global
        for (int i = lid; i < tlen * HEAD_DIM; i += tt) {
            int s = i / HEAD_DIM;
            int d_dim = i % HEAD_DIM;
            int kv_idx = kv0 + s;
            int g_off = kv_base + kv_idx * p.kv_stride_l + d_dim * p.kv_stride_d;
            sK[i] = p.k[g_off];
            sV[i] = p.v[g_off];
        }
        __syncthreads();
        int lim = tlen;
        if (p.causal_offset >= 0 && q_row < p.q_len) {
            int ep = q_row + p.causal_offset + 1;
            if (kv0 >= ep)
                lim = 0;
            else if (kv0 + tlen > ep)
                lim = ep - kv0;
        }
        int mask_row_base = mask_batch_base + q_row * p.mask_q_stride;
        for (int s = 0; s < lim; s++) {
            const bf16* kr = sK + s * HEAD_DIM + gpos * DPT;
            float part = 0.0f;
 #pragma unroll
            for (int i = 0; i < DPT; i += 8) {
                float k8[8];
                ld8(kr + i, k8);
 #pragma unroll
                for (int j = 0; j < 8; j++)
                    part = fmaf(qreg[i + j], k8[j], part);
            }
            float dot = group_reduce_sum<G>(part, gmask);
            int kv_idx = kv0 + s;
            if (p.use_mask && p.mask && !p.mask[mask_row_base + kv_idx])
                dot = -FLT_MAX;
            float nm = fmaxf(m, dot);
            float al = __expf(m - nm);
            float be = __expf(dot - nm);
            l = l * al + be;
            const bf16* vr = sV + s * HEAD_DIM + gpos * DPT;
 #pragma unroll
            for (int i = 0; i < DPT; i += 8) {
                float v8[8];
                ld8(vr + i, v8);
 #pragma unroll
                for (int j = 0; j < 8; j++)
                    acc[i + j] = fmaf(v8[j], be, acc[i + j] * al);
            }
            m = nm;
        }
        __syncthreads();
    }
    if (q_row < p.q_len) {
        // O: stride-based write
        int o_off = batch * p.q_stride_b + q_head * p.q_stride_h
                  + q_row * p.q_stride_l + gpos * DPT * p.q_stride_d;
        float rl = (l > 1e-10f) ? (1.0f / l) : 0.0f;
 #pragma unroll
        for (int i = 0; i < DPT; i++)
            p.o[o_off + i * p.q_stride_d] = __float2bfloat16(acc[i] * rl);
    }
 }
@@ -0,0 +1,196 @@
 #pragma once
 #include <cfloat>
 #include <cuda_bf16.h>
 #include "attn_common.h"
 #include "attn_mma_utils.cuh"
 using bf16 = __nv_bfloat16;
 // Tensor-core prefill flash attention (raw mma.sync PTX).
 // One warp owns BR=16 query rows. S = Q@K^T and O = P@V run on bf16 tensor
 // cores via mma.sync.m16n8k16 (f32 accumulate). Q fragments are loaded once
 // straight from global into the mma A-operand layout (no smem staging) and
 // kept resident in registers across the tile loop. S, O, and the online-softmax
 // stats (m, l) also live in registers.
 // Shared memory is statically sized via template parameters — no dynamic
 // allocation. The mma fragment layout is used directly: the S accumulator
 // (f32) maps element-for-element onto the P matrix_a (bf16) operand, so
 // softmax needs no shuffle repack; row reductions fold across the 4-lane
 // thread group. Templated on <HEAD_DIM, WARPS, BC> with BC a multiple of 16.
 //
 // Software pipeline: K/V are double-buffered and loaded via cp.async one tile
 // ahead, so the next tile streams from global memory while the current tile's
 // tensor-core math runs — hiding load latency (long_scoreboard). A single
 // __syncthreads per tile both publishes the freshly loaded tile cross-warp and
 // (because it runs before the next prefetch) guards the buffer being refilled,
 // so no second barrier is needed. Predicated cp.async (cp_async_16_pred)
 // zero-fills rows past kv_len, unifying full and partial tiles on one path.
 // BC=32 (D<=128) amortizes the per-tile wait+barrier+loop overhead over more
 // tensor-core work — this kernel is latency-bound (low occupancy from high
 // register pressure), so fewer, larger tiles beat many tiny ones.
 //
 // Optimizations: load Q fragments directly from global in mma A-operand layout
 // (no sQ staging, no prologue barriers); post-multiply scale in float after
 // S=Q@K^T to avoid bf16 precision loss; packed bf16x2 output stores;
 // causal tile skipping (block-level prefetch bound + warp-level compute skip);
 // XOR swizzle (swiz_col) → eliminates ldmatrix bank conflicts without LD
 // padding (LD=HEAD_DIM).
 template <int HEAD_DIM, int WARPS, int BC>
 __global__ void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
    constexpr int BR = 16;
    constexpr int KD = HEAD_DIM / 16;  // Q/K k-tiles
    constexpr int NC8 = BC / 8;        // S n-tiles (N=8 each)
    constexpr int KT2 = BC / 16;       // P k-tiles (K=16 each)
    constexpr int DN8 = HEAD_DIM / 8;  // O n-tiles (N=8 each)
    constexpr int LD = HEAD_DIM;   // XOR swizzle (swiz_col) handles bank conflicts
    constexpr int SWIZ_MASK = (HEAD_DIM >= 64) ? 7 : (HEAD_DIM / 8 - 1);  // chunk bits, stay within LD
    const int warp = threadIdx.x / 32;
    const int lane = threadIdx.x % 32;
    const int gid = lane >> 2;   // 0..7  → rows gid, gid+8
    const int tid4 = lane & 3;   // 0..3
    const int nthreads = WARPS * 32;
    const int q_head = blockIdx.y;
    const int batch = blockIdx.z;
    const int kv_head = q_head / (p.q_head / p.kv_head);
    const int qrow0 = (blockIdx.x * WARPS + warp) * BR;
    // ---- Static shared memory: double-buffered K/V ----
    // K/V are double-buffered (STAGES=2): the next tile's cp.async load runs
    // while the current tile's tensor-core math executes, hiding global-load
    // latency (FA2-style software pipeline). No dynamic smem / carveout opt-in.
    constexpr int STAGES = 2;
    __shared__ __align__(16) bf16 sK[STAGES * BC * LD];
    __shared__ __align__(16) bf16 sV[STAGES * BC * LD];
    // Load Q fragments straight from global into mma A-operand layout.
    // stride_row = p.q_stride_l for prefill (multi-q rows across q_len).
    // See attn_mma_utils.cuh for the shared template.
    const int q_base = batch * p.q_stride_b + q_head * p.q_stride_h;
    const int qra = qrow0 + gid;
    const int qrb = qrow0 + gid + 8;
    const bool va = qra < p.q_len, vb = qrb < p.q_len;
    unsigned Qa[KD][4];
    load_q_mma_frags<KD>(p.q + q_base, p.q_stride_l, p.q_stride_d,
                          qra, qrb, va, vb, tid4, Qa);
    float Oacc[DN8][4];
 #pragma unroll
    for (int j = 0; j < DN8; j++)
        Oacc[j][0] = Oacc[j][1] = Oacc[j][2] = Oacc[j][3] = 0.0f;
    float m0 = -FLT_MAX, m1 = -FLT_MAX, l0 = 0.0f, l1 = 0.0f;
    // KV: stride-based base
    const int kv_base = batch * p.kv_stride_b + kv_head * p.kv_stride_h;
    const int tiles = (p.kv_len + BC - 1) / BC;
    const int qr0 = qrow0 + gid;      // row for c0/c1
    const int qr1 = qrow0 + gid + 8;  // row for c2/c3
    // Causal tile-skip bounds (no-op when causal_offset < 0)
    const int use_skip = (p.causal_offset >= 0) ? 1 : 0;
    const int max_kv = qrow0 + BR - 1 + p.causal_offset;
    const int block_max_kv =
        blockIdx.x * WARPS * BR + WARPS * BR - 1 + p.causal_offset;
    const int has_mask = p.use_mask && p.mask;
    // Last active tile: block-level causal bound (all warps in the block share
    // the K/V load, so the prefetch range is the block max, not per-warp).
    int t_end = tiles - 1;
    if (use_skip) {
        int bt = block_max_kv / BC;
        if (bt < t_end) t_end = bt;
    }
    constexpr int VEC = 8;  // bf16 per cp.async unit (16 bytes)
    constexpr int TOTAL = BC * HEAD_DIM;
    // ---- Load tile lambda: predicated cp.async ----
    // Issue cp.async loads for tile `ti` into shared buffer `buf`. Predicated
    // loads zero-fill rows past kv_len, so partial tiles need no scalar path.
    auto load_tile = [&](int ti, int buf) {
        int kv0 = ti * BC;
        bf16* dK = sK + buf * BC * LD;
        bf16* dV = sV + buf * BC * LD;
 #pragma unroll
        for (int i = threadIdx.x * VEC; i < TOTAL; i += nthreads * VEC) {
            int r = i / HEAD_DIM, d = i % HEAD_DIM;
            int kc = kv0 + r;
            bool valid = kc < p.kv_len;
            int off = r * LD + swiz_col(d, r, SWIZ_MASK);
            int g_off = kv_base + kc * p.kv_stride_l + d * p.kv_stride_d;
            cp_async_16_pred(&dK[off], &p.k[g_off], valid);
            cp_async_16_pred(&dV[off], &p.v[g_off], valid);
        }
        cp_async_commit();
    };
    // ---- Prologue: issue first tile load ----
    load_tile(0, 0);
    for (int ti = 0; ti <= t_end; ti++) {
        int buf = ti & 1;
        // Wait for the current tile's async copies, then a single barrier: it
        // both publishes this tile's data cross-warp AND guarantees the prior
        // compute on the buffer we are about to refill has finished. Issuing
        // the next tile's load *after* this barrier lets one barrier cover both
        // hazards (vs two), while the load still overlaps this tile's math.
        cp_async_wait_group<0>();
        __syncthreads();
        if (ti < t_end) load_tile(ti + 1, (ti + 1) & 1);
        const bf16* bK = sK + buf * BC * LD;
        const bf16* bV = sV + buf * BC * LD;
        int kv0 = ti * BC;
        // Warp-level causal skip
        if (!use_skip || kv0 <= max_kv) {
        // S = Q @ K^T + scale + online softmax + O += P @ V
        float Sacc[NC8][4];
        mma_compute_scores<KD, NC8>(Qa, bK, LD, SWIZ_MASK, lane, Sacc);
        // post-multiply scale in float (no bf16 precision loss from pre-scaling Q)
        #pragma unroll
        for (int n8 = 0; n8 < NC8; n8++)
            Sacc[n8][0] *= p.scale, Sacc[n8][1] *= p.scale,
            Sacc[n8][2] *= p.scale, Sacc[n8][3] *= p.scale;
        int maxc0 = (p.causal_offset >= 0) ? min(p.kv_len, qr0 + p.causal_offset + 1)
                                        : p.kv_len;
        int maxc1 = (p.causal_offset >= 0) ? min(p.kv_len, qr1 + p.causal_offset + 1)
                                        : p.kv_len;
        mma_softmax_tile<NC8, DN8>(kv0, maxc0, maxc1,
                                    qr0, qr1,
                                    p.mask_b_stride, p.mask_q_stride,
                                    batch,
                                    p.mask, has_mask,
                                    Sacc, Oacc, m0, m1, l0, l1, lane);
        mma_pv_accumulate<DN8, KT2>(Sacc, bV, LD, SWIZ_MASK, lane, Oacc);
        }  // if active (warp-level causal skip)
    }
    // ---- write output ---- (packed bf16x2 stores: one 32-bit STG per pair,
    // halves store count and removes the uncoalesced scalar-store penalty)
    float rl0 = (l0 > 1e-20f) ? (1.0f / l0) : 0.0f;
    float rl1 = (l1 > 1e-20f) ? (1.0f / l1) : 0.0f;
    // O: stride-based write
    const int o_base = batch * p.q_stride_b + q_head * p.q_stride_h;
 #pragma unroll
    for (int dn8 = 0; dn8 < DN8; dn8++) {
        int d = dn8 * 8 + 2 * tid4;
        if (qr0 < p.q_len) {
            __nv_bfloat162 v = __floats2bfloat162_rn(Oacc[dn8][0] * rl0,
                                                     Oacc[dn8][1] * rl0);
            *reinterpret_cast<__nv_bfloat162*>(&p.o[o_base + qr0 * p.q_stride_l + d * p.q_stride_d]) = v;
        }
        if (qr1 < p.q_len) {
            __nv_bfloat162 v = __floats2bfloat162_rn(Oacc[dn8][2] * rl1,
                                                     Oacc[dn8][3] * rl1);
            *reinterpret_cast<__nv_bfloat162*>(&p.o[o_base + qr1 * p.q_stride_l + d * p.q_stride_d]) = v;
        }
    }
 }
--- a/Show More
+++ b/Show More
		`@@ -0,0 +1,2 @@`
							`# Source directory for CUDA kernels — build-time only.`
							`# Compiled .so files live in astrAI/_ext/.`