refactor: simplify inference engine and backend dispatch

- merge _generate_streaming/_generate_non_streaming into single _generate() with stream flag - delete dead GenerationRequest class and generate_with_request method - inline _next_token helper into generate_async - replace flash-attn double-checked locking with functools.lru_cache - extract _write_and_gather_kv helper shared by TorchNative/FlashAttn backends - inline _kv_cache_is_contiguous into its sole call site in FlashAttnBackend - change default backend priority from flash>cuda>torch to cuda>flash>torch - add ASTR_BACKEND env var to override default backend at resolve time - add supports_graph() static method to AttentionBackend ABC, override in CudaBackend - replace isinstance(get_backend(), CudaBackend) with get_backend().supports_graph() in executor - add torch.cuda.is_available() guard to CudaBackend.supports()
feat: add timed() context manager and backend supports()
2026-08-07 22:28:48 +08:00 · 2026-08-07 20:51:30 +08:00 · 2026-08-07 20:04:10 +08:00 · 2026-08-07 19:45:59 +08:00 · 2026-08-07 18:32:08 +08:00 · 2026-08-07 18:21:02 +08:00
187 changed files with 28190 additions and 6189 deletions
@@ -4,6 +4,8 @@
 # Allow necessary files
 !astrai/
 !scripts/
-!assets/
+!docs/
 !csrc/
 !setup.py
 !pyproject.toml
 !README.md
@@ -0,0 +1,100 @@
 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
          if-no-files-found: error
  build-cuda-linux:
    name: Build CUDA wheel (Linux, ${{ matrix.cuda_tag }})
    runs-on: ubuntu-latest
    strategy:
      fail-fast: false
      matrix:
        include:
          - cuda_tag: "cu128"
            cuda_ver: "12.8.0"
          - cuda_tag: "cu130"
            cuda_ver: "13.0.0"
    steps:
      - uses: actions/checkout@v4
      - uses: actions/setup-python@v5
        with:
          python-version: "3.12"
      - name: Install torch (${{ matrix.cuda_tag }})
        run: |
          pip install torch --index-url https://download.pytorch.org/whl/${{ matrix.cuda_tag }}
      - name: Setup CUDA (${{ matrix.cuda_ver }})
        uses: Jimver/cuda-toolkit@v0.2.35
        with:
          cuda: "${{ matrix.cuda_ver }}"
      - 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-${{ matrix.cuda_tag }}
          path: dist/*.whl
          if-no-files-found: error
  release:
    name: Attach wheels to release
    needs: [build-pure, build-cuda-linux]
    runs-on: ubuntu-latest
    permissions:
      contents: write
    steps:
      - name: Download pure-Python wheel
        uses: actions/download-artifact@v4
        with:
          name: pure-wheel
          path: release-assets/pure
      - name: Download CUDA wheels (all variants)
        uses: actions/download-artifact@v4
        with:
          pattern: cuda-wheel-linux-*
          merge-multiple: true
          path: release-assets/cuda
      - name: Verify release assets
        shell: bash
        run: |
          set -euo pipefail
          pure_wheels=(release-assets/pure/*.whl)
          cuda_wheels=(release-assets/cuda/*.whl)
          test "${#pure_wheels[@]}" -eq 1
          test "${#cuda_wheels[@]}" -ge 1
      - name: Create release & upload assets
        uses: softprops/action-gh-release@v2
        with:
          files: |
            release-assets/pure/*.whl
            release-assets/cuda/*.whl
          tag_name: ${{ github.ref_name }}
          generate_release_notes: true
@@ -5,8 +5,17 @@
 !*/
 # Allow specific file types and root files
-!*.py
+!astrai/**/*.py
-!*.sh
+!scripts/**/*.py
 !tests/**/*.py
 !csrc/**/*.py
 !csrc/CMakeLists.txt
 !csrc/**/*.cu
 !csrc/**/*.h
 !csrc/**/*.cuh
 !scripts/**/*.sh
 # Allow GitHub files
 !/.github/**
@@ -16,8 +25,13 @@
 !/.dockerignore
 !/Dockerfile
 !/docker-compose.yml
-!/assets/**
+!/docs/**
 !/CONTRIBUTING.md
 !/LICENSE
 !/pyproject.toml
-!/README.md
+!/README.md
 # Allow extension modules (only source .py)
 !/astrai/extension/**/*.py
 # Allow build files
 !/setup.py
@@ -5,7 +5,7 @@ Thank you for your interest in contributing! This document provides step-by-step
 ## Quick Start
 ```bash
-git clone https://github.com/your-username/AstrAI.git
+git clone https://github.com/ViperEkura/AstrAI.git
 cd AstrAI
 pip install -e ".[dev]"     # install with dev dependencies (pytest, ruff)
 ```
@@ -20,9 +20,6 @@ Run the following checks **in order** — CI will reject if any fail.
 ruff format .
 ```
 > **Note**: `ruff format` may rename parameters (e.g. `mask` → `attn_mask`).  
 > Always review the diff after formatting.
 ### 2. Import sorting
 ```bash
@@ -44,7 +41,7 @@ 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
+### 4. (Optional) Full pre-commit check script
 If you have Git Bash available:
@@ -52,12 +49,17 @@ If you have Git Bash available:
 bash scripts/pre_commit.sh
 ```
-This runs format check, import sort check, and tests in one go.
+The script installs development dependencies by default, then runs the format
 check, import sort check, and tests. If dependencies are already installed, use:
 ```bash
 bash scripts/pre_commit.sh --skip-deps
 ```
 ## Commit Style
 ```
-fix/feat/chore/docs/refactor/perf/test/style/ci/build/revert : short description (~50 chars)
+type: short description (~50 chars)
 - bullet point body (each ~60 chars)
 ```
@@ -73,7 +75,7 @@ fix/feat/chore/docs/refactor/perf/test/style/ci/build/revert : short description
 |---------|-------|-----|
 | `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` |
+| Pre-commit check script fails | Dependency install, tests, or lint failed | Fix the failing step; use `--skip-deps` only when dependencies are already installed |
 | Tests fail with tempdir left | Test crash | Clean `%TEMP%` manually |
 ## Submitting Changes
@@ -93,7 +95,7 @@ fix/feat/chore/docs/refactor/perf/test/style/ci/build/revert : short description
 ## License
-By contributing, you agree that your contributions will be licensed under the [GPL-3.0 License](LICENSE).
+By contributing, you agree that your contributions will be licensed under the [Apache-2.0 License](LICENSE).
 ---
@@ -1,8 +1,16 @@
 # AstrAI Dockerfile - Multi-stage Build (Optimized)
 #
 # CUDA version selection:
 #   docker build -t astrai .
 #   docker build -t astrai --build-arg CUDA_TAG=cu128 .
 #   docker build -t astrai --build-arg CUDA_TAG=cu130 .
 # Default: cu128
 # Build stage - use base image with minimal build tools
 FROM ubuntu:24.04 AS builder
 ARG CUDA_TAG=cu128
 WORKDIR /app
 # Install Python 3.12 and minimal build dependencies
@@ -20,10 +28,12 @@ ENV PATH="/opt/venv/bin:$PATH"
 # Copy source code and install (deps read from pyproject.toml)
 COPY astrai/ ./astrai/
 COPY csrc/ ./csrc/
 COPY setup.py .
 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/${CUDA_TAG}"
 # Production stage
 FROM ubuntu:24.04 AS production
@@ -43,7 +53,7 @@ ENV PATH="/opt/venv/bin:$PATH"
 # Copy application code
 COPY astrai/ ./astrai/
 COPY scripts/ ./scripts/
-COPY assets/ ./assets/
+COPY docs/ ./docs/
 COPY pyproject.toml .
 COPY README.md .
@@ -1,674 +1,201 @@
-                    GNU GENERAL PUBLIC LICENSE
+                                 Apache License
-                       Version 3, 29 June 2007
+                           Version 2.0, January 2004
-
+                        http://www.apache.org/licenses/
- Copyright (C) 2007 Free Software Foundation, Inc. <https://fsf.org/>
+
- Everyone is permitted to copy and distribute verbatim copies
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
- of this license document, but changing it is not allowed.
+
-
+   1. Definitions.
-                            Preamble
+
-
+      "License" shall mean the terms and conditions for use, reproduction,
-  The GNU General Public License is a free, copyleft license for
+      and distribution as defined by Sections 1 through 9 of this document.
-software and other kinds of works.
+
-
+      "Licensor" shall mean the copyright owner or entity authorized by
-  The licenses for most software and other practical works are designed
+      the copyright owner that is granting the License.
-to take away your freedom to share and change the works.  By contrast,
+
-the GNU General Public License is intended to guarantee your freedom to
+      "Legal Entity" shall mean the union of the acting entity and all
-share and change all versions of a program--to make sure it remains free
+      other entities that control, are controlled by, or are under common
-software for all its users.  We, the Free Software Foundation, use the
+      control with that entity. For the purposes of this definition,
-GNU General Public License for most of our software; it applies also to
+      "control" means (i) the power, direct or indirect, to cause the
-any other work released this way by its authors.  You can apply it to
+      direction or management of such entity, whether by contract or
-your programs, too.
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
-
+      outstanding shares, or (iii) beneficial ownership of such entity.
-  When we speak of free software, we are referring to freedom, not
+
-price.  Our General Public Licenses are designed to make sure that you
+      "You" (or "Your") shall mean an individual or Legal Entity
-have the freedom to distribute copies of free software (and charge for
+      exercising permissions granted by this License.
-them if you wish), that you receive source code or can get it if you
+
-want it, that you can change the software or use pieces of it in new
+      "Source" form shall mean the preferred form for making modifications,
-free programs, and that you know you can do these things.
+      including but not limited to software source code, documentation
-
+      source, and configuration files.
-  To protect your rights, we need to prevent others from denying you
+
-these rights or asking you to surrender the rights.  Therefore, you have
+      "Object" form shall mean any form resulting from mechanical
-certain responsibilities if you distribute copies of the software, or if
+      transformation or translation of a Source form, including but
-you modify it: responsibilities to respect the freedom of others.
+      not limited to compiled object code, generated documentation,
-
+      and conversions to other media types.
-  For example, if you distribute copies of such a program, whether
+
-gratis or for a fee, you must pass on to the recipients the same
+      "Work" shall mean the work of authorship, whether in Source or
-freedoms that you received.  You must make sure that they, too, receive
+      Object form, made available under the License, as indicated by a
-or can get the source code.  And you must show them these terms so they
+      copyright notice that is included in or attached to the work
-know their rights.
+      (an example is provided in the Appendix below).
-
+
-  Developers that use the GNU GPL protect your rights with two steps:
+      "Derivative Works" shall mean any work, whether in Source or Object
-(1) assert copyright on the software, and (2) offer you this License
+      form, that is based on (or derived from) the Work and for which the
-giving you legal permission to copy, distribute and/or modify it.
+      editorial revisions, annotations, elaborations, or other modifications
-
+      represent, as a whole, an original work of authorship. For the purposes
-  For the developers' and authors' protection, the GPL clearly explains
+      of this License, Derivative Works shall not include works that remain
-that there is no warranty for this free software.  For both users' and
+      separable from, or merely link (or bind by name) to the interfaces of,
-authors' sake, the GPL requires that modified versions be marked as
+      the Work and Derivative Works thereof.
-changed, so that their problems will not be attributed erroneously to
+
-authors of previous versions.
+      "Contribution" shall mean any work of authorship, including
-
+      the original version of the Work and any modifications or additions
-  Some devices are designed to deny users access to install or run
+      to that Work or Derivative Works thereof, that is intentionally
-modified versions of the software inside them, although the manufacturer
+      submitted to Licensor for inclusion in the Work by the copyright owner
-can do so.  This is fundamentally incompatible with the aim of
+      or by an individual or Legal Entity authorized to submit on behalf of
-protecting users' freedom to change the software.  The systematic
+      the copyright owner. For the purposes of this definition, "submitted"
-pattern of such abuse occurs in the area of products for individuals to
+      means any form of electronic, verbal, or written communication sent
-use, which is precisely where it is most unacceptable.  Therefore, we
+      to the Licensor or its representatives, including but not limited to
-have designed this version of the GPL to prohibit the practice for those
+      communication on electronic mailing lists, source code control systems,
-products.  If such problems arise substantially in other domains, we
+      and issue tracking systems that are managed by, or on behalf of, the
-stand ready to extend this provision to those domains in future versions
+      Licensor for the purpose of discussing and improving the Work, but
-of the GPL, as needed to protect the freedom of users.
+      excluding communication that is conspicuously marked or otherwise
-
+      designated in writing by the copyright owner as "Not a Contribution."
-  Finally, every program is threatened constantly by software patents.
+
-States should not allow patents to restrict development and use of
+      "Contributor" shall mean Licensor and any individual or Legal Entity
-software on general-purpose computers, but in those that do, we wish to
+      on behalf of whom a Contribution has been received by Licensor and
-avoid the special danger that patents applied to a free program could
+      subsequently incorporated within the Work.
-make it effectively proprietary.  To prevent this, the GPL assures that
+
-patents cannot be used to render the program non-free.
+   2. Grant of Copyright License. Subject to the terms and conditions of
-
+      this License, each Contributor hereby grants to You a perpetual,
-  The precise terms and conditions for copying, distribution and
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
-modification follow.
+      copyright license to reproduce, prepare Derivative Works of,
-
+      publicly display, publicly perform, sublicense, and distribute the
-                       TERMS AND CONDITIONS
+      Work and such Derivative Works in Source or Object form.
-
+
-  0. Definitions.
+   3. Grant of Patent License. Subject to the terms and conditions of
-
+      this License, each Contributor hereby grants to You a perpetual,
-  "This License" refers to version 3 of the GNU General Public License.
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
-
+      (except as stated in this section) patent license to make, have made,
-  "Copyright" also means copyright-like laws that apply to other kinds of
+      use, offer to sell, sell, import, and otherwise transfer the Work,
-works, such as semiconductor masks.
+      where such license applies only to those patent claims licensable
-
+      by such Contributor that are necessarily infringed by their
-  "The Program" refers to any copyrightable work licensed under this
+      Contribution(s) alone or by combination of their Contribution(s)
-License.  Each licensee is addressed as "you".  "Licensees" and
+      with the Work to which such Contribution(s) was submitted. If You
-"recipients" may be individuals or organizations.
+      institute patent litigation against any entity (including a
-
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
-  To "modify" a work means to copy from or adapt all or part of the work
+      or a Contribution incorporated within the Work constitutes direct
-in a fashion requiring copyright permission, other than the making of an
+      or contributory patent infringement, then any patent licenses
-exact copy.  The resulting work is called a "modified version" of the
+      granted to You under this License for that Work shall terminate
-earlier work or a work "based on" the earlier work.
+      as of the date such litigation is filed.
-
+
-  A "covered work" means either the unmodified Program or a work based
+   4. Redistribution. You may reproduce and distribute copies of the
-on the Program.
+      Work or Derivative Works thereof in any medium, with or without
-
+      modifications, and in Source or Object form, provided that You
-  To "propagate" a work means to do anything with it that, without
+      meet the following conditions:
-permission, would make you directly or secondarily liable for
+
-infringement under applicable copyright law, except executing it on a
+      (a) You must give any other recipients of the Work or
-computer or modifying a private copy.  Propagation includes copying,
+          Derivative Works a copy of this License; and
-distribution (with or without modification), making available to the
+
-public, and in some countries other activities as well.
+      (b) You must cause any modified files to carry prominent notices
-
+          stating that You changed the files; and
-  To "convey" a work means any kind of propagation that enables other
+
-parties to make or receive copies.  Mere interaction with a user through
+      (c) You must retain, in the Source form of any Derivative Works
-a computer network, with no transfer of a copy, is not conveying.
+          that You distribute, all copyright, patent, trademark, and
-
+          attribution notices from the Source form of the Work,
-  An interactive user interface displays "Appropriate Legal Notices"
+          excluding those notices that do not pertain to any part of
-to the extent that it includes a convenient and prominently visible
+          the Derivative Works; and
-feature that (1) displays an appropriate copyright notice, and (2)
+
-tells the user that there is no warranty for the work (except to the
+      (d) If the Work includes a "NOTICE" text file as part of its
-extent that warranties are provided), that licensees may convey the
+          distribution, then any Derivative Works that You distribute must
-work under this License, and how to view a copy of this License.  If
+          include a readable copy of the attribution notices contained
-the interface presents a list of user commands or options, such as a
+          within such NOTICE file, excluding those notices that do not
-menu, a prominent item in the list meets this criterion.
+          pertain to any part of the Derivative Works, in at least one
-
+          of the following places: within a NOTICE text file distributed
-  1. Source Code.
+          as part of the Derivative Works; within the Source form or
-
+          documentation, if provided along with the Derivative Works; or,
-  The "source code" for a work means the preferred form of the work
+          within a display generated by the Derivative Works, if and
-for making modifications to it.  "Object code" means any non-source
+          wherever such third-party notices normally appear. The contents
-form of a work.
+          of the NOTICE file are for informational purposes only and
-
+          do not modify the License. You may add Your own attribution
-  A "Standard Interface" means an interface that either is an official
+          notices within Derivative Works that You distribute, alongside
-standard defined by a recognized standards body, or, in the case of
+          or as an addendum to the NOTICE text from the Work, provided
-interfaces specified for a particular programming language, one that
+          that such additional attribution notices cannot be construed
-is widely used among developers working in that language.
+          as modifying the License.
-
+
-  The "System Libraries" of an executable work include anything, other
+      You may add Your own copyright statement to Your modifications and
-than the work as a whole, that (a) is included in the normal form of
+      may provide additional or different license terms and conditions
-packaging a Major Component, but which is not part of that Major
+      for use, reproduction, or distribution of Your modifications, or
-Component, and (b) serves only to enable use of the work with that
+      for any such Derivative Works as a whole, provided Your use,
-Major Component, or to implement a Standard Interface for which an
+      reproduction, and distribution of the Work otherwise complies with
-implementation is available to the public in source code form.  A
+      the conditions stated in this License.
-"Major Component", in this context, means a major essential component
+
-(kernel, window system, and so on) of the specific operating system
+   5. Submission of Contributions. Unless You explicitly state otherwise,
-(if any) on which the executable work runs, or a compiler used to
+      any Contribution intentionally submitted for inclusion in the Work
-produce the work, or an object code interpreter used to run it.
+      by You to the Licensor shall be under the terms and conditions of
-
+      this License, without any additional terms or conditions.
-  The "Corresponding Source" for a work in object code form means all
+      Notwithstanding the above, nothing herein shall supersede or modify
-the source code needed to generate, install, and (for an executable
+      the terms of any separate license agreement you may have executed
-work) run the object code and to modify the work, including scripts to
+      with Licensor regarding such Contributions.
-control those activities.  However, it does not include the work's
+
-System Libraries, or general-purpose tools or generally available free
+   6. Trademarks. This License does not grant permission to use the trade
-programs which are used unmodified in performing those activities but
+      names, trademarks, service marks, or product names of the Licensor,
-which are not part of the work.  For example, Corresponding Source
+      except as required for reasonable and customary use in describing the
-includes interface definition files associated with source files for
+      origin of the Work and reproducing the content of the NOTICE file.
-the work, and the source code for shared libraries and dynamically
+
-linked subprograms that the work is specifically designed to require,
+   7. Disclaimer of Warranty. Unless required by applicable law or
-such as by intimate data communication or control flow between those
+      agreed to in writing, Licensor provides the Work (and each
-subprograms and other parts of the work.
+      Contributor provides its Contributions) on an "AS IS" BASIS,
-
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
-  The Corresponding Source need not include anything that users
+      implied, including, without limitation, any warranties or conditions
-can regenerate automatically from other parts of the Corresponding
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
-Source.
+      PARTICULAR PURPOSE. You are solely responsible for determining the
-
+      appropriateness of using or redistributing the Work and assume any
-  The Corresponding Source for a work in source code form is that
+      risks associated with Your exercise of permissions under this License.
-same work.
+
-
+   8. Limitation of Liability. In no event and under no legal theory,
-  2. Basic Permissions.
+      whether in tort (including negligence), contract, or otherwise,
-
+      unless required by applicable law (such as deliberate and grossly
-  All rights granted under this License are granted for the term of
+      negligent acts) or agreed to in writing, shall any Contributor be
-copyright on the Program, and are irrevocable provided the stated
+      liable to You for damages, including any direct, indirect, special,
-conditions are met.  This License explicitly affirms your unlimited
+      incidental, or consequential damages of any character arising as a
-permission to run the unmodified Program.  The output from running a
+      result of this License or out of the use or inability to use the
-covered work is covered by this License only if the output, given its
+      Work (including but not limited to damages for loss of goodwill,
-content, constitutes a covered work.  This License acknowledges your
+      work stoppage, computer failure or malfunction, or any and all
-rights of fair use or other equivalent, as provided by copyright law.
+      other commercial damages or losses), even if such Contributor
-
+      has been advised of the possibility of such damages.
-  You may make, run and propagate covered works that you do not
+
-convey, without conditions so long as your license otherwise remains
+   9. Accepting Warranty or Additional Liability. While redistributing
-in force.  You may convey covered works to others for the sole purpose
+      the Work or Derivative Works thereof, You may choose to offer,
-of having them make modifications exclusively for you, or provide you
+      and charge a fee for, acceptance of support, warranty, indemnity,
-with facilities for running those works, provided that you comply with
+      or other liability obligations and/or rights consistent with this
-the terms of this License in conveying all material for which you do
+      License. However, in accepting such obligations, You may act only
-not control copyright.  Those thus making or running the covered works
+      on Your own behalf and on Your sole responsibility, not on behalf
-for you must do so exclusively on your behalf, under your direction
+      of any other Contributor, and only if You agree to indemnify,
-and control, on terms that prohibit them from making any copies of
+      defend, and hold each Contributor harmless for any liability
-your copyrighted material outside their relationship with you.
+      incurred by, or claims asserted against, such Contributor by reason
-
+      of your accepting any such warranty or additional liability.
-  Conveying under any other circumstances is permitted solely under
+
-the conditions stated below.  Sublicensing is not allowed; section 10
+   END OF TERMS AND CONDITIONS
-makes it unnecessary.
+
-
+   APPENDIX: How to apply the Apache License to your work.
-  3. Protecting Users' Legal Rights From Anti-Circumvention Law.
+
-
+      To apply the Apache License to your work, attach the following
-  No covered work shall be deemed part of an effective technological
+      boilerplate notice, with the fields enclosed by brackets "[]"
-measure under any applicable law fulfilling obligations under article
+      replaced with your own identifying information. (Don't include
-11 of the WIPO copyright treaty adopted on 20 December 1996, or
+      the brackets!)  The text should be enclosed in the appropriate
-similar laws prohibiting or restricting circumvention of such
+      comment syntax for the file format. We also recommend that a
-measures.
+      file or class name and description of purpose be included on the
-
+      same "printed page" as the copyright notice for easier
-  When you convey a covered work, you waive any legal power to forbid
+      identification within third-party archives.
-circumvention of technological measures to the extent such circumvention
+
-is effected by exercising rights under this License with respect to
+   Copyright [yyyy] [name of copyright owner]
-the covered work, and you disclaim any intention to limit operation or
+
-modification of the work as a means of enforcing, against the work's
+   Licensed under the Apache License, Version 2.0 (the "License");
-users, your or third parties' legal rights to forbid circumvention of
+   you may not use this file except in compliance with the License.
-technological measures.
+   You may obtain a copy of the License at
-
+
-  4. Conveying Verbatim Copies.
+       http://www.apache.org/licenses/LICENSE-2.0
-
+
-  You may convey verbatim copies of the Program's source code as you
+   Unless required by applicable law or agreed to in writing, software
-receive it, in any medium, provided that you conspicuously and
+   distributed under the License is distributed on an "AS IS" BASIS,
-appropriately publish on each copy an appropriate copyright notice;
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-keep intact all notices stating that this License and any
+   See the License for the specific language governing permissions and
-non-permissive terms added in accord with section 7 apply to the code;
+   limitations under the License.
 keep intact all notices of the absence of any warranty; and give all
 recipients a copy of this License along with the Program.
  You may charge any price or no price for each copy that you convey,
 and you may offer support or warranty protection for a fee.
  5. Conveying Modified Source Versions.
  You may convey a work based on the Program, or the modifications to
 produce it from the Program, in the form of source code under the
 terms of section 4, provided that you also meet all of these conditions:
    a) The work must carry prominent notices stating that you modified
    it, and giving a relevant date.
    b) The work must carry prominent notices stating that it is
    released under this License and any conditions added under section
    7.  This requirement modifies the requirement in section 4 to
    "keep intact all notices".
    c) You must license the entire work, as a whole, under this
    License to anyone who comes into possession of a copy.  This
    License will therefore apply, along with any applicable section 7
    additional terms, to the whole of the work, and all its parts,
    regardless of how they are packaged.  This License gives no
    permission to license the work in any other way, but it does not
    invalidate such permission if you have separately received it.
    d) If the work has interactive user interfaces, each must display
    Appropriate Legal Notices; however, if the Program has interactive
    interfaces that do not display Appropriate Legal Notices, your
    work need not make them do so.
  A compilation of a covered work with other separate and independent
 works, which are not by their nature extensions of the covered work,
 and which are not combined with it such as to form a larger program,
 in or on a volume of a storage or distribution medium, is called an
 "aggregate" if the compilation and its resulting copyright are not
 used to limit the access or legal rights of the compilation's users
 beyond what the individual works permit.  Inclusion of a covered work
 in an aggregate does not cause this License to apply to the other
 parts of the aggregate.
  6. Conveying Non-Source Forms.
  You may convey a covered work in object code form under the terms
 of sections 4 and 5, provided that you also convey the
 machine-readable Corresponding Source under the terms of this License,
 in one of these ways:
    a) Convey the object code in, or embodied in, a physical product
    (including a physical distribution medium), accompanied by the
    Corresponding Source fixed on a durable physical medium
    customarily used for software interchange.
    b) Convey the object code in, or embodied in, a physical product
    (including a physical distribution medium), accompanied by a
    written offer, valid for at least three years and valid for as
    long as you offer spare parts or customer support for that product
    model, to give anyone who possesses the object code either (1) a
    copy of the Corresponding Source for all the software in the
    product that is covered by this License, on a durable physical
    medium customarily used for software interchange, for a price no
    more than your reasonable cost of physically performing this
    conveying of source, or (2) access to copy the
    Corresponding Source from a network server at no charge.
    c) Convey individual copies of the object code with a copy of the
    written offer to provide the Corresponding Source.  This
    alternative is allowed only occasionally and noncommercially, and
    only if you received the object code with such an offer, in accord
    with subsection 6b.
    d) Convey the object code by offering access from a designated
    place (gratis or for a charge), and offer equivalent access to the
    Corresponding Source in the same way through the same place at no
    further charge.  You need not require recipients to copy the
    Corresponding Source along with the object code.  If the place to
    copy the object code is a network server, the Corresponding Source
    may be on a different server (operated by you or a third party)
    that supports equivalent copying facilities, provided you maintain
    clear directions next to the object code saying where to find the
    Corresponding Source.  Regardless of what server hosts the
    Corresponding Source, you remain obligated to ensure that it is
    available for as long as needed to satisfy these requirements.
    e) Convey the object code using peer-to-peer transmission, provided
    you inform other peers where the object code and Corresponding
    Source of the work are being offered to the general public at no
    charge under subsection 6d.
  A separable portion of the object code, whose source code is excluded
 from the Corresponding Source as a System Library, need not be
 included in conveying the object code work.
  A "User Product" is either (1) a "consumer product", which means any
 tangible personal property which is normally used for personal, family,
 or household purposes, or (2) anything designed or sold for incorporation
 into a dwelling.  In determining whether a product is a consumer product,
 doubtful cases shall be resolved in favor of coverage.  For a particular
 product received by a particular user, "normally used" refers to a
 typical or common use of that class of product, regardless of the status
 of the particular user or of the way in which the particular user
 actually uses, or expects or is expected to use, the product.  A product
 is a consumer product regardless of whether the product has substantial
 commercial, industrial or non-consumer uses, unless such uses represent
 the only significant mode of use of the product.
  "Installation Information" for a User Product means any methods,
 procedures, authorization keys, or other information required to install
 and execute modified versions of a covered work in that User Product from
 a modified version of its Corresponding Source.  The information must
 suffice to ensure that the continued functioning of the modified object
 code is in no case prevented or interfered with solely because
 modification has been made.
  If you convey an object code work under this section in, or with, or
 specifically for use in, a User Product, and the conveying occurs as
 part of a transaction in which the right of possession and use of the
 User Product is transferred to the recipient in perpetuity or for a
 fixed term (regardless of how the transaction is characterized), the
 Corresponding Source conveyed under this section must be accompanied
 by the Installation Information.  But this requirement does not apply
 if neither you nor any third party retains the ability to install
 modified object code on the User Product (for example, the work has
 been installed in ROM).
  The requirement to provide Installation Information does not include a
 requirement to continue to provide support service, warranty, or updates
 for a work that has been modified or installed by the recipient, or for
 the User Product in which it has been modified or installed.  Access to a
 network may be denied when the modification itself materially and
 adversely affects the operation of the network or violates the rules and
 protocols for communication across the network.
  Corresponding Source conveyed, and Installation Information provided,
 in accord with this section must be in a format that is publicly
 documented (and with an implementation available to the public in
 source code form), and must require no special password or key for
 unpacking, reading or copying.
  7. Additional Terms.
  "Additional permissions" are terms that supplement the terms of this
 License by making exceptions from one or more of its conditions.
 Additional permissions that are applicable to the entire Program shall
 be treated as though they were included in this License, to the extent
 that they are valid under applicable law.  If additional permissions
 apply only to part of the Program, that part may be used separately
 under those permissions, but the entire Program remains governed by
 this License without regard to the additional permissions.
  When you convey a copy of a covered work, you may at your option
 remove any additional permissions from that copy, or from any part of
 it.  (Additional permissions may be written to require their own
 removal in certain cases when you modify the work.)  You may place
 additional permissions on material, added by you to a covered work,
 for which you have or can give appropriate copyright permission.
  Notwithstanding any other provision of this License, for material you
 add to a covered work, you may (if authorized by the copyright holders of
 that material) supplement the terms of this License with terms:
    a) Disclaiming warranty or limiting liability differently from the
    terms of sections 15 and 16 of this License; or
    b) Requiring preservation of specified reasonable legal notices or
    author attributions in that material or in the Appropriate Legal
    Notices displayed by works containing it; or
    c) Prohibiting misrepresentation of the origin of that material, or
    requiring that modified versions of such material be marked in
    reasonable ways as different from the original version; or
    d) Limiting the use for publicity purposes of names of licensors or
    authors of the material; or
    e) Declining to grant rights under trademark law for use of some
    trade names, trademarks, or service marks; or
    f) Requiring indemnification of licensors and authors of that
    material by anyone who conveys the material (or modified versions of
    it) with contractual assumptions of liability to the recipient, for
    any liability that these contractual assumptions directly impose on
    those licensors and authors.
  All other non-permissive additional terms are considered "further
 restrictions" within the meaning of section 10.  If the Program as you
 received it, or any part of it, contains a notice stating that it is
 governed by this License along with a term that is a further
 restriction, you may remove that term.  If a license document contains
 a further restriction but permits relicensing or conveying under this
 License, you may add to a covered work material governed by the terms
 of that license document, provided that the further restriction does
 not survive such relicensing or conveying.
  If you add terms to a covered work in accord with this section, you
 must place, in the relevant source files, a statement of the
 additional terms that apply to those files, or a notice indicating
 where to find the applicable terms.
  Additional terms, permissive or non-permissive, may be stated in the
 form of a separately written license, or stated as exceptions;
 the above requirements apply either way.
  8. Termination.
  You may not propagate or modify a covered work except as expressly
 provided under this License.  Any attempt otherwise to propagate or
 modify it is void, and will automatically terminate your rights under
 this License (including any patent licenses granted under the third
 paragraph of section 11).
  However, if you cease all violation of this License, then your
 license from a particular copyright holder is reinstated (a)
 provisionally, unless and until the copyright holder explicitly and
 finally terminates your license, and (b) permanently, if the copyright
 holder fails to notify you of the violation by some reasonable means
 prior to 60 days after the cessation.
  Moreover, your license from a particular copyright holder is
 reinstated permanently if the copyright holder notifies you of the
 violation by some reasonable means, this is the first time you have
 received notice of violation of this License (for any work) from that
 copyright holder, and you cure the violation prior to 30 days after
 your receipt of the notice.
  Termination of your rights under this section does not terminate the
 licenses of parties who have received copies or rights from you under
 this License.  If your rights have been terminated and not permanently
 reinstated, you do not qualify to receive new licenses for the same
 material under section 10.
  9. Acceptance Not Required for Having Copies.
  You are not required to accept this License in order to receive or
 run a copy of the Program.  Ancillary propagation of a covered work
 occurring solely as a consequence of using peer-to-peer transmission
 to receive a copy likewise does not require acceptance.  However,
 nothing other than this License grants you permission to propagate or
 modify any covered work.  These actions infringe copyright if you do
 not accept this License.  Therefore, by modifying or propagating a
 covered work, you indicate your acceptance of this License to do so.
  10. Automatic Licensing of Downstream Recipients.
  Each time you convey a covered work, the recipient automatically
 receives a license from the original licensors, to run, modify and
 propagate that work, subject to this License.  You are not responsible
 for enforcing compliance by third parties with this License.
  An "entity transaction" is a transaction transferring control of an
 organization, or substantially all assets of one, or subdividing an
 organization, or merging organizations.  If propagation of a covered
 work results from an entity transaction, each party to that
 transaction who receives a copy of the work also receives whatever
 licenses to the work the party's predecessor in interest had or could
 give under the previous paragraph, plus a right to possession of the
 Corresponding Source of the work from the predecessor in interest, if
 the predecessor has it or can get it with reasonable efforts.
  You may not impose any further restrictions on the exercise of the
 rights granted or affirmed under this License.  For example, you may
 not impose a license fee, royalty, or other charge for exercise of
 rights granted under this License, and you may not initiate litigation
 (including a cross-claim or counterclaim in a lawsuit) alleging that
 any patent claim is infringed by making, using, selling, offering for
 sale, or importing the Program or any portion of it.
  11. Patents.
  A "contributor" is a copyright holder who authorizes use under this
 License of the Program or a work on which the Program is based.  The
 work thus licensed is called the contributor's "contributor version".
  A contributor's "essential patent claims" are all patent claims
 owned or controlled by the contributor, whether already acquired or
 hereafter acquired, that would be infringed by some manner, permitted
 by this License, of making, using, or selling its contributor version,
 but do not include claims that would be infringed only as a
 consequence of further modification of the contributor version.  For
 purposes of this definition, "control" includes the right to grant
 patent sublicenses in a manner consistent with the requirements of
 this License.
  Each contributor grants you a non-exclusive, worldwide, royalty-free
 patent license under the contributor's essential patent claims, to
 make, use, sell, offer for sale, import and otherwise run, modify and
 propagate the contents of its contributor version.
  In the following three paragraphs, a "patent license" is any express
 agreement or commitment, however denominated, not to enforce a patent
 (such as an express permission to practice a patent or covenant not to
 sue for patent infringement).  To "grant" such a patent license to a
 party means to make such an agreement or commitment not to enforce a
 patent against the party.
  If you convey a covered work, knowingly relying on a patent license,
 and the Corresponding Source of the work is not available for anyone
 to copy, free of charge and under the terms of this License, through a
 publicly available network server or other readily accessible means,
 then you must either (1) cause the Corresponding Source to be so
 available, or (2) arrange to deprive yourself of the benefit of the
 patent license for this particular work, or (3) arrange, in a manner
 consistent with the requirements of this License, to extend the patent
 license to downstream recipients.  "Knowingly relying" means you have
 actual knowledge that, but for the patent license, your conveying the
 covered work in a country, or your recipient's use of the covered work
 in a country, would infringe one or more identifiable patents in that
 country that you have reason to believe are valid.
  If, pursuant to or in connection with a single transaction or
 arrangement, you convey, or propagate by procuring conveyance of, a
 covered work, and grant a patent license to some of the parties
 receiving the covered work authorizing them to use, propagate, modify
 or convey a specific copy of the covered work, then the patent license
 you grant is automatically extended to all recipients of the covered
 work and works based on it.
  A patent license is "discriminatory" if it does not include within
 the scope of its coverage, prohibits the exercise of, or is
 conditioned on the non-exercise of one or more of the rights that are
 specifically granted under this License.  You may not convey a covered
 work if you are a party to an arrangement with a third party that is
 in the business of distributing software, under which you make payment
 to the third party based on the extent of your activity of conveying
 the work, and under which the third party grants, to any of the
 parties who would receive the covered work from you, a discriminatory
 patent license (a) in connection with copies of the covered work
 conveyed by you (or copies made from those copies), or (b) primarily
 for and in connection with specific products or compilations that
 contain the covered work, unless you entered into that arrangement,
 or that patent license was granted, prior to 28 March 2007.
  Nothing in this License shall be construed as excluding or limiting
 any implied license or other defenses to infringement that may
 otherwise be available to you under applicable patent law.
  12. No Surrender of Others' Freedom.
  If conditions are imposed on you (whether by court order, agreement or
 otherwise) that contradict the conditions of this License, they do not
 excuse you from the conditions of this License.  If you cannot convey a
 covered work so as to satisfy simultaneously your obligations under this
 License and any other pertinent obligations, then as a consequence you may
 not convey it at all.  For example, if you agree to terms that obligate you
 to collect a royalty for further conveying from those to whom you convey
 the Program, the only way you could satisfy both those terms and this
 License would be to refrain entirely from conveying the Program.
  13. Use with the GNU Affero General Public License.
  Notwithstanding any other provision of this License, you have
 permission to link or combine any covered work with a work licensed
 under version 3 of the GNU Affero General Public License into a single
 combined work, and to convey the resulting work.  The terms of this
 License will continue to apply to the part which is the covered work,
 but the special requirements of the GNU Affero General Public License,
 section 13, concerning interaction through a network will apply to the
 combination as such.
  14. Revised Versions of this License.
  The Free Software Foundation may publish revised and/or new versions of
 the GNU General Public License from time to time.  Such new versions will
 be similar in spirit to the present version, but may differ in detail to
 address new problems or concerns.
  Each version is given a distinguishing version number.  If the
 Program specifies that a certain numbered version of the GNU General
 Public License "or any later version" applies to it, you have the
 option of following the terms and conditions either of that numbered
 version or of any later version published by the Free Software
 Foundation.  If the Program does not specify a version number of the
 GNU General Public License, you may choose any version ever published
 by the Free Software Foundation.
  If the Program specifies that a proxy can decide which future
 versions of the GNU General Public License can be used, that proxy's
 public statement of acceptance of a version permanently authorizes you
 to choose that version for the Program.
  Later license versions may give you additional or different
 permissions.  However, no additional obligations are imposed on any
 author or copyright holder as a result of your choosing to follow a
 later version.
  15. Disclaimer of Warranty.
  THERE IS NO WARRANTY FOR THE PROGRAM, TO THE EXTENT PERMITTED BY
 APPLICABLE LAW.  EXCEPT WHEN OTHERWISE STATED IN WRITING THE COPYRIGHT
 HOLDERS AND/OR OTHER PARTIES PROVIDE THE PROGRAM "AS IS" WITHOUT WARRANTY
 OF ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING, BUT NOT LIMITED TO,
 THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
 PURPOSE.  THE ENTIRE RISK AS TO THE QUALITY AND PERFORMANCE OF THE PROGRAM
 IS WITH YOU.  SHOULD THE PROGRAM PROVE DEFECTIVE, YOU ASSUME THE COST OF
 ALL NECESSARY SERVICING, REPAIR OR CORRECTION.
  16. Limitation of Liability.
  IN NO EVENT UNLESS REQUIRED BY APPLICABLE LAW OR AGREED TO IN WRITING
 WILL ANY COPYRIGHT HOLDER, OR ANY OTHER PARTY WHO MODIFIES AND/OR CONVEYS
 THE PROGRAM AS PERMITTED ABOVE, BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY
 GENERAL, SPECIAL, INCIDENTAL OR CONSEQUENTIAL DAMAGES ARISING OUT OF THE
 USE OR INABILITY TO USE THE PROGRAM (INCLUDING BUT NOT LIMITED TO LOSS OF
 DATA OR DATA BEING RENDERED INACCURATE OR LOSSES SUSTAINED BY YOU OR THIRD
 PARTIES OR A FAILURE OF THE PROGRAM TO OPERATE WITH ANY OTHER PROGRAMS),
 EVEN IF SUCH HOLDER OR OTHER PARTY HAS BEEN ADVISED OF THE POSSIBILITY OF
 SUCH DAMAGES.
  17. Interpretation of Sections 15 and 16.
  If the disclaimer of warranty and limitation of liability provided
 above cannot be given local legal effect according to their terms,
 reviewing courts shall apply local law that most closely approximates
 an absolute waiver of all civil liability in connection with the
 Program, unless a warranty or assumption of liability accompanies a
 copy of the Program in return for a fee.
                     END OF TERMS AND CONDITIONS
            How to Apply These Terms to Your New Programs
  If you develop a new program, and you want it to be of the greatest
 possible use to the public, the best way to achieve this is to make it
 free software which everyone can redistribute and change under these terms.
  To do so, attach the following notices to the program.  It is safest
 to attach them to the start of each source file to most effectively
 state the exclusion of warranty; and each file should have at least
 the "copyright" line and a pointer to where the full notice is found.
    <one line to give the program's name and a brief idea of what it does.>
    Copyright (C) <year>  <name of author>
    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.
    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.
    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <https://www.gnu.org/licenses/>.
 Also add information on how to contact you by electronic and paper mail.
  If the program does terminal interaction, make it output a short
 notice like this when it starts in an interactive mode:
    <program>  Copyright (C) <year>  <name of author>
    This program comes with ABSOLUTELY NO WARRANTY; for details type `show w'.
    This is free software, and you are welcome to redistribute it
    under certain conditions; type `show c' for details.
 The hypothetical commands `show w' and `show c' should show the appropriate
 parts of the General Public License.  Of course, your program's commands
 might be different; for a GUI interface, you would use an "about box".
  You should also get your employer (if you work as a programmer) or school,
 if any, to sign a "copyright disclaimer" for the program, if necessary.
 For more information on this, and how to apply and follow the GNU GPL, see
 <https://www.gnu.org/licenses/>.
  The GNU General Public License does not permit incorporating your program
 into proprietary programs.  If your program is a subroutine library, you
 may consider it more useful to permit linking proprietary applications with
 the library.  If this is what you want to do, use the GNU Lesser General
 Public License instead of this License.  But first, please read
 <https://www.gnu.org/licenses/why-not-lgpl.html>.
@@ -1,6 +1,6 @@
 <div align="center">
-  <img src="assets/images/logo.png" width="auto" alt="Logo">
+  <img src="docs/images/logo.png" width="auto" alt="Logo">
  <p>
    <strong>A lightweight Transformer training & inference framework</strong>
  </p>
@@ -8,27 +8,28 @@
 <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/badge/license-Apache--2.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>
 <div align="center">
  <a href="#english">English</a> •
-  <a href="assets/docs/README-zh-CN.md">中文</a> •
+  <a href="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>
 ## 📖 Table of Contents
- [Features](#features)
+- [Overview](#overview)
- [Quick Start](#quick-start)
+- [Getting Started](#getting-started)
 - [Demo](#demo)
 - [Documentation](#documentation)
 - [Contributing](#contributing)
 - [Community](#community)
@@ -39,49 +40,67 @@
 <a id="english"></a>
 ## English
-### Features
+### Overview
- 🚀 **High Performance**: Optimized for both training and inference with efficient parallelization.
+AstrAI is an end-to-end Transformer framework for building, training, evaluating, and serving models. It provides a compact PyTorch codebase for the complete model lifecycle, from declarative data preprocessing and distributed training to continuous-batching inference and OpenAI/Anthropic-compatible APIs.
 - 🔧 **Flexible**: Support for seq/sft/dpo/grpo training, customizable model architectures.
 - 💡 **Easy to Use**: Simple API with comprehensive examples and demos.
 - 📦 **Lightweight**: Minimal dependencies, easy to deploy.
 - 🔬 **Research‑Friendly**: Modular design, easy to experiment with new ideas.
 - 🤗 **HuggingFace-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
+| Area | Capabilities |
 |---|---|
 | **Models** | Autoregressive language models and embedding models with GQA, MLA, MoE, RoPE, and extensible attention/FFN components |
 | **Training** | Pre-training (`seq`), supervised fine-tuning (`sft`), DPO, and GRPO with gradient accumulation, checkpointing, DDP, and FSDP |
 | **Data** | Declarative JSON preprocessing, configurable masking and packing, binary/JSONL storage, and streaming datasets |
 | **Inference** | Continuous batching, paged KV cache, radix prefix caching, streaming generation, and Torch/CUDA/FlashAttention backends |
 | **Serving** | FastAPI server with OpenAI and Anthropic chat completion protocols, including SSE streaming and tool calls |
 | **Evaluation** | Perplexity, MMLU, HumanEval, IFEval, IFD, and ROUGE evaluation tools |
 | **Extensibility** | Factory and registry architecture for models, datasets, training strategies, callbacks, kernels, and protocol components |
-#### Installation
+### Getting Started
 End-to-end walkthrough in 5 steps:
 **1. Install**
 AstrAI requires Python 3.12+ and pins PyTorch exactly to `2.11.0`. Training, `scripts/tools/generate.py`, generation evaluations, and the generation demos require CUDA; CPU support is limited to components with an explicit CPU device path, such as the HTTP server and direct-scoring evaluations.
 ```bash
 git clone https://github.com/ViperEkura/AstrAI.git
 cd AstrAI
-pip install -e .
+pip install -e .                                          # pure PyTorch (no CUDA kernels)
 # CSRC_KERNELS=true pip install -e . --no-build-isolation  # optional: fused CUDA kernels
 # pip install -e ".[dev]"                                  # dev dependencies (pytest, ruff)
 ```
-For development dependencies:
+**2. Download model**
 ```bash
-pip install -e ".[dev]"
+python scripts/demo/download.py    # downloads 1B checkpoint to params/
 ```
-#### Download Pre-trained Model
+**3. Preprocess data**
-Download pre-trained model weights (1B bilingual checkpoint) to `params/`:
+Create `pretrain.json` (preprocessing config for `seq` strategy):
 ```json
 {
    "version": 1,
    "input": {"sections": [{"field": "text", "action": "train"}]},
    "preprocessing": {"max_seq_len": 2048},
    "output": {"storage_format": "bin"}
 }
 ```
 ```bash
-python scripts/demo/download.py
+python scripts/tools/preprocess.py data/*.jsonl -o output/ -c pretrain.json
 ```
-Or download manually from [HuggingFace](https://huggingface.co/ViperEk/KHAOSZ) into `params/`.
+**4. Train**
 #### Train a Model
 ```bash
 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 \
@@ -90,9 +109,7 @@ nohup python scripts/tools/train.py \
    --warmup_ratio=0.05 \
    --max_lr=1e-4 \
    --max_grad_norm=1.0 \
-    --adamw_beta1=0.9 \
+    --weight_decay=0.1 \
    --adamw_beta2=0.95 \
    --adamw_weight_decay=0.01 \
    --window_size=2048 \
    --ckpt_interval=10000 \
    --ckpt_dir=./checkpoint \
@@ -101,15 +118,54 @@ nohup python scripts/tools/train.py \
    > out.log 2> err.log &
 ```
-Full reference at [Parameter Guide](assets/docs/params.md).
+**5. Serve & query**
-#### Generate Text
+```bash
 # Terminal 1: start server
 python scripts/tools/server.py --param_path ./params --device cuda
 # Terminal 2: query
 curl http://localhost:8000/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{"messages":[{"role":"user","content":"Hello"}],"max_tokens":512}'
 ```
 ### Demo
 Check out the demos in the `scripts/demo/` folder:
 ```bash
 # Download model weights (required before running demos)
 python scripts/demo/download.py                      # model → params/
 # Single-turn interactive streaming prompt loop (no conversation 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
@@ -123,9 +179,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
@@ -136,94 +189,53 @@ docker run --gpus all -v /path/to/data:/data -it astrai:latest
 # Docker Compose (GPU, default)
 docker compose up -d
-# Docker Compose (CPU only)
+# Docker Compose CPU server profile (CUDA-only generation scripts/demos are unavailable)
 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](docs/guides/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/BV1fuLB6yEj6).
 ### Documentation
 | Document | Description |
 |----------|-------------|
-| [Parameter Guide](./assets/docs/params.md) | Training & inference parameters |
+| [Get Started](./docs/get-started.md) | Installation and quickstart |
-| [Architecture](./assets/docs/architecture.md) | System architecture, class diagram & design patterns |
+| [CLI Reference](./docs/guides/params.md) | Parameters for all CLI tools (train, server, generate, preprocess) |
-| [Training](./assets/docs/training.md) | Training loop, strategies & formulas |
+| [Preprocessing](./docs/guides/preprocessing.md) | Declarative JSON-driven data preprocessing |
-| [Inference](./assets/docs/inference.md) | KVCache, continuous batching, sampling & HTTP API |
+| [Training](./docs/guides/training.md) | Training loop, strategies & formulas |
-| [Data Flow](./assets/docs/dataflow.md) | Data pipeline, storage backends & dataset architecture |
+| [Inference](./docs/guides/inference.md) | KVCache, continuous batching, sampling & HTTP API |
 | [Evaluation](./docs/guides/evaluation.md) | HumanEval, MMLU, PPL, ROUGE, IFD, IFEval |
 | [Distributed](./docs/guides/distributed.md) | Multi-GPU DDP / FSDP training |
 | [Architecture](./docs/developer/architecture.md) | System architecture, class diagram & design patterns |
 | [Data Flow](./docs/developer/dataflow.md) | Data pipeline, storage backends & dataset architecture |
 | [Internals](./docs/developer/internals.md) | Training internals: loss formulas, callback lifecycle, KV cache |
 | [CUDA Kernels](./docs/developer/cuda_kernels.md) | Custom CUDA attention kernels & benchmarks |
 ### Contributing
@@ -240,14 +252,14 @@ 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
-This project is licensed under the [GPL-3.0 License](LICENSE).
+This project is licensed under the [Apache-2.0 License](LICENSE).
 ---
 <div align="center">
  <em>A lightweight Transformer framework designed for both high performance and ease of use.</em>
-</div>
+</div>
@@ -1,259 +0,0 @@
 <div align="center">
  <img src="../images/logo.png" width="auto" alt="Logo">
  <div>
    <a href="../../README.md">English</a> • 
    <a href="#chinese">中文</a>
  </div>
  <p>
    <strong>轻量级 Transformer 训练与推理框架</strong>
  </p>
 </div>
 <div align="center">
  <img src="https://img.shields.io/badge/python-3.12+-blue.svg" alt="python">
  <img src="https://img.shields.io/badge/license-GPL--3.0-blue.svg" alt="license">
  <img src="https://img.shields.io/github/v/release/ViperEkura/AstrAI?color=76bad9" alt="release">
  <img src="https://img.shields.io/badge/dynamic/json?url=https%3A%2F%2Fapi.github.com%2Frepos%2FViperEkura%2FAstrAI&query=%24.stargazers_count&label=stars&suffix=%20stars&color=76bad9" alt="stars">
  <img src="https://img.shields.io/badge/dynamic/json?url=https%3A%2F%2Fapi.github.com%2Frepos%2FViperEkura%2FAstrAI&query=%24.forks_count&label=forks&suffix=%20forks&color=76bad9" alt="forks">
 </div>
 <br>
 <div align="center">
  <a href="../../README.md">English</a> •
  <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>
 </div>
 <br>
 ## 📖 目录
 - [特性](#特性)
 - [快速开始](#快速开始)
 - [文档](#文档)
 - [贡献](#贡献)
 - [社区](#社区)
 - [许可证](#许可证)
 ---
 <a id="chinese"></a>
 ## 中文
 ### 特性
 - 🚀 **高性能**: 训练与推理双向优化，高效并行。
 - 🔧 **灵活**: 支持 seq/sft/dpo/grpo 多种训练方式，可定制模型架构。
 - 💡 **易用**: 简洁的 API 与丰富的示例、演示。
 - 📦 **轻量**: 依赖少，部署简单。
 - 🔬 **研究友好**: 模块化设计，便于实验新想法。
 - 🤗 **HuggingFace 风格 API**: 类 HuggingFace 的 AutoModel/AutoTokenizer 接口，方便加载模型和分词器。
 - 🔌 **双 API 兼容**: 同时支持 OpenAI 和 Anthropic 聊天补全 API，开箱即用。
 ### 快速开始
 #### 安装
 ```bash
 git clone https://github.com/ViperEkura/AstrAI.git
 cd AstrAI
 pip install -e .
 ```
 安装开发依赖：
 ```bash
 pip install -e ".[dev]"
 ```
 #### 下载预训练模型
 下载预训练模型权重（1B 双语检查点）到 `params/` 目录：
 ```bash
 python scripts/demo/download.py
 ```
 或从 [HuggingFace](https://huggingface.co/ViperEk/KHAOSZ) 手动下载放入 `params/`。
 #### 训练模型
 ```bash
 export CUDA_VISIBLE_DEVICES=0,1,2,3
 nohup python scripts/tools/train.py \
    --nprocs=4 \
    --train_type=seq \
    --data_root_path=/path/to/dataset \
    --param_path=/path/to/model \
    --batch_per_device=4 \
    --grad_accum_steps=8 \
    --warmup_ratio=0.05 \
    --max_lr=1e-4 \
    --max_grad_norm=1.0 \
    --adamw_beta1=0.9 \
    --adamw_beta2=0.95 \
    --adamw_weight_decay=0.01 \
    --window_size=2048 \
    --ckpt_interval=10000 \
    --ckpt_dir=./checkpoint \
    --random_seed=3407 \
    --label_smoothing=0.05 \
    > out.log 2> err.log &
 ```
 完整参数列表见[参数说明](./params.md)。
 #### 文本生成
 ```bash
 python scripts/tools/generate.py \
    --param_path /path/to/model \
    --input_json_file /path/to/input.json \
    --output_json_file /path/to/output.json
 ```
 #### Docker
 使用 Docker 构建和运行（推荐用于 GPU 环境）：
 ```bash
 # 构建镜像
 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
 # 挂载数据卷
 docker run --gpus all -v /path/to/data:/data -it astrai:latest
 # Docker Compose（GPU，默认）
 docker compose up -d
 # Docker Compose（仅 CPU）
 docker compose --profile cpu up -d
 ```
 > **注意**: 必须使用 `--gpus all` 才能启用 CUDA 支持，否则 `torch.cuda.is_available()` 将返回 `False`。
 #### 启动 HTTP 服务
 启动推理服务器，支持 OpenAI 和 Anthropic 兼容的 HTTP API：
 ```bash
 python -m scripts.tools.server --port 8000 --device cuda
 ```
 发起请求：
 ```bash
 # OpenAI 兼容
 curl -X POST http://localhost:8000/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{
    "messages": [{"role": "user", "content": "你好"}],
    "max_tokens": 512
  }'
 # OpenAI 兼容流式
 curl -X POST http://localhost:8000/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{
    "messages": [{"role": "user", "content": "讲个故事"}],
    "stream": true,
    "max_tokens": 500
  }'
 # Anthropic 兼容
 curl -X POST http://localhost:8000/v1/messages \
  -H "Content-Type: application/json" \
  -d '{
    "model": "astrai",
    "system": "你是一个乐于助人的助手。",
    "messages": [{"role": "user", "content": "你好"}],
    "max_tokens": 512
  }'
 # Anthropic 兼容流式并设置停止序列
 curl -X POST http://localhost:8000/v1/messages \
  -H "Content-Type: application/json" \
  -d '{
    "model": "astrai",
    "messages": [{"role": "user", "content": "写个故事"}],
    "max_tokens": 500,
    "stream": true,
    "stop_sequences": ["结束"]
  }'
 # 健康检查
 curl http://localhost:8000/health
 ```
 #### 演示
 查看 `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/BV1fuLB6yEj6) 上的视频演示。
 ### 文档
 | 文档 | 说明 |
 |------|------|
 | [参数说明](./params.md) | 训练与推理参数配置 |
 | [架构文档](./architecture.md) | 系统架构、类图与设计模式 |
 | [训练文档](./training.md) | 训练循环、策略与公式 |
 | [推理文档](./inference.md) | KVCache、连续批处理、采样与 HTTP API |
 | [数据流程](./dataflow.md) | 数据管道、存储后端与数据集架构 |
 ### 贡献
 我们欢迎贡献！请参阅[贡献指南](../../CONTRIBUTING.md)了解详情。
 1. Fork 本仓库。
 2. 创建功能分支。
 3. 提交更改。
 4. 发起 Pull Request。
 重大更改请先开 issue 讨论。
 ### 社区
 - **GitHub Issues**: [问题追踪](https://github.com/ViperEkura/AstrAI/issues)
 - **Discussions**: [GitHub 讨论区](https://github.com/ViperEkura/AstrAI/discussions)
 - **HuggingFace**: [模型中心](https://huggingface.co/ViperEk)
 ### 许可证
 本项目采用 [GPL-3.0 许可证](../../LICENSE)。
 ---
 <div align="center">
  <em>专为高性能与易用性设计的轻量级 Transformer 框架。</em>
 </div>
@@ -1,57 +0,0 @@
 # Data Flow
 This document describes the data pipeline: from raw text to model input tensors.
 ## Overview
 ```
 Raw Text → AutoTokenizer → Token IDs → .h5/.bin → Dataset → Sampler → DataLoader → Training/Inference
 ```
 ## Data Preparation
 Raw text is tokenized via `AutoTokenizer.encode()` and saved as HDF5 (`.h5`) or binary (`.bin` + `meta.json`) files with keyed tensor groups.
 Storage format is auto-detected by `detect_format()`; backends are dispatched via registry:
 ```
 StoreFactory.create("h5")  → H5Store
 StoreFactory.create("bin") → MmapStore
 ```
 H5 backend supports shared memory via `.share_memory_()`. Bin (mmap) uses OS page-cache sharing natively.
 ## Data Keys by Training Type
 | Type | Storage Keys |
 |------|-------------|
 | `seq` | `sequence` (→ input_ids, target_ids via offset-by-1) |
 | `sft` | `sequence`, `loss_mask` |
 | `dpo` | `chosen`, `rejected`, `chosen_mask`, `rejected_mask` |
 | `grpo` | `prompts`, `responses`, `masks`, `rewards` |
 ## Dataset Architecture
 ```
 DatasetFactory.load(train_type, load_path, window_size, stride, storage_type)
  → StoreFactory.create(detect_format(path))
    → Store._data[Dict[str, List[Tensor]]] + _cum[Dict[str, List[int]]]
      → BaseDataset.__getitem__(idx)
        → sliding window [begin, end) via get_index(idx)
 ```
 `window_size` = max input length, `stride` = step between consecutive samples (defaults to `window_size`).
 ## Sampler
 `ResumableDistributedSampler` supports checkpoint-aware distributed sampling:
 - Tracks `start_epoch` / `start_iter` for resume
 - Shuffle via `torch.Generator(seed + epoch)`
 - Per-replica index slicing for DDP
 ## DataLoader
 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-05-28
@@ -1,148 +0,0 @@
 # Inference
 ## KV Cache
 At decode time, only the last query token matters. All previous K/V are cached to avoid recomputation:
 $$
 o_n = \sum_j \text{softmax}\left(\frac{q_n k_j}{\sqrt{d_k}}\right) v_j
 $$
 RoPE is applied **before** KV cache write, not after — otherwise position encoding drift occurs.
 ## KVCache System
 Six classes working together:
 ```
 KVCache (facade)
  ├── PagePool         orchestrates page allocation + prefix matching
  │     ├── Allocator   bitmask-based page allocator + ref-count + LRU eviction (inside PagePool)
  │     └── PrefixCache hash-based prefix matching (page_hash via polynomial hash) (inside PagePool)
  ├── TaskTable        maps task_id → page_table + cached token count
  ├── Storage          k_cache / v_cache tensors (n_layers × n_pages × page_size × n_kv_heads × head_dim)
  └── KvcacheView      bundles Storage + page_table + total_len for attention layers (returned by bind())
 ```
 `KVCache.bind(page_table, total_len)` returns a `KvcacheView` used by attention layers via `write()` / `gather()`.
 ## Continuous Batching
 `InferenceScheduler` runs a daemon thread with a 4-phase loop:
 ```
 1. Cleanup → Remove finished tasks, free KV pages
 2. Refill  → Pop from waiting_queue, task_alloc pages, activate
 3. Prefill → Group by (prompt_len, start_pos), run full forward
 4. Decode  → Pick largest same-position group, single-token forward
 ```
 ## Sampling (Strategy Pattern)
 ```
 BaseSamplingStrategy (ABC)
  ├── TemperatureStrategy
  ├── TopKStrategy
  └── TopPStrategy
 ```
 `SamplingPipeline` 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
 ```
 `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) |
 | `temperature` | float | 1.0 | Sampling temperature (>= 0.0) |
 | `top_p` | float | 1.0 | Nucleus threshold |
 | `top_k` | int | 50 | Top-k count |
 | `max_tokens` | Optional[int] | None | Max generation length |
 | `stream` | bool | False | Stream output |
 ## Engine API
 ```python
 # Non-streaming
 engine.generate("Hello", stream=False)          # -> str
 engine.generate(["A", "B"], stream=False)       # -> List[str]
 # Streaming
 engine.generate("Hello", stream=True)           # -> Generator[str]
 engine.generate(["A", "B"], stream=True)        # -> Generator[Tuple[int, str]]
 # Async
 await engine.generate_async("Hello", ...)       # -> AsyncGenerator[str]
 ```
 > Document Update Time: 2026-05-28
@@ -1,99 +0,0 @@
 # Parameter Documentation
 ## Training Parameters
 ### Basic Parameters
 | Parameter | Description | Default |
 |-----------|-------------|---------|
 | `--train_type` | Training type (`seq`, `sft`, `dpo`, `grpo`) | required |
 | `--data_root_path` | Dataset root directory | required |
 | `--param_path` | Model parameters or checkpoint path | required |
 | `--n_epoch` | Total training epochs | 1 |
 | `--batch_per_device` | Batch size per device | 1 |
 | `--grad_accum_steps` | Gradient accumulation steps between optimizer steps | 1 |
 ### Learning Rate Scheduling
 | Parameter | Description | Default |
 |-----------|-------------|---------|
 | `--warmup_ratio` | Fraction of total steps used for LR warmup | 0.05 |
 | `--max_lr` | Maximum learning rate (cosine decay after warmup) | 3e-4 |
 | `--max_grad_norm` | Maximum gradient norm for clipping | 1.0 |
 ### Optimizer (AdamW)
 | Parameter | Description | Default |
 |-----------|-------------|---------|
 | `--adamw_beta1` | AdamW beta1 | 0.9 |
 | `--adamw_beta2` | AdamW beta2 | 0.95 |
 | `--adamw_weight_decay` | AdamW weight decay | 0.01 |
 ### Data Loading
 | Parameter | Description | Default |
 |-----------|-------------|---------|
 | `--window_size` | Max input sequence length | model config `max_len` |
 | `--stride` | Stride for sliding window over sequences | None |
 | `--random_seed` | Random seed for reproducibility | 3407 |
 | `--num_workers` | DataLoader worker processes | 4 |
 | `--no_pin_memory` | Disable pin_memory (enabled by default) | (flag) |
 ### Checkpoint & Resume
 | Parameter | Description | Default |
 |-----------|-------------|---------|
 | `--ckpt_interval` | Iterations between checkpoints | 5000 |
 | `--ckpt_dir` | Checkpoint save directory | checkpoint |
 | `--start_epoch` | Resume from epoch (0 = from scratch) | 0 |
 | `--start_batch` | Resume from batch iteration | 0 |
 ### Distributed Training
 | Parameter | Description | Default |
 |-----------|-------------|---------|
 | `--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 |
 ### Strategy-specific
 | Parameter | Description | Default | Used by |
 |-----------|-------------|---------|---------|
 | `--dpo_beta` | DPO beta value | 0.1 | `dpo` |
 | `--label_smoothing` | Label smoothing for cross-entropy loss | 0.05 | `seq`, `sft` |
 | `--group_size` | GRPO group size | 4 | `grpo` |
 | `--grpo_clip_eps` | GRPO clipping epsilon | 0.2 | `grpo` |
 | `--grpo_kl_coef` | GRPO KL penalty coefficient | 0.01 | `grpo` |
 | `--grpo_sync_interval` | GRPO ref_model sync interval (steps) | 200 | `grpo` |
 ### Usage Example
 ```bash
 export CUDA_VISIBLE_DEVICES=0,1,2,3
 nohup python scripts/tools/train.py \
    --nprocs=4 \
    --train_type=seq \
    --data_root_path=/path/to/dataset \
    --param_path=/path/to/model \
    --batch_per_device=4 \
    --grad_accum_steps=8 \
    --warmup_ratio=0.05 \
    --max_lr=1e-4 \
    --max_grad_norm=1.0 \
    --adamw_beta1=0.9 \
    --adamw_beta2=0.95 \
    --adamw_weight_decay=0.01 \
    --window_size=2048 \
    --ckpt_interval=10000 \
    --ckpt_dir=./checkpoint \
    --random_seed=3407 \
    --label_smoothing=0.05 \
    > out.log 2> err.log &
 ```
 ---
 > Document Update Time: 2026-05-24
@@ -1,227 +0,0 @@
 # Training
 ## Model Architecture
 The model uses a decoder-only Transformer with **GQA** (Grouped Query Attention) and optional **MLA** (Multi-head Latent Attention). 1.0 billion parameters, Chinese–English bilingual.
 ```mermaid
 flowchart TB
    subgraph Layers["Transformer Layers"]
        direction TB
        A[Input Embedding] --> B[Transformer Block\nLayer 1]
        B --> C[Transformer Block\nLayer ...]
        C --> D[Transformer Block\nLayer ...]
        D --> E[RMSNorm]
        E --> F[Linear]
        F --> G[SoftMax]
    end
    subgraph TransformerBlock["Transformer Block"]
        direction TB
        H[x] --> I[RMSNorm]
        I --> J[Linear → Q/K/V]
        J --> K[Q]; J --> L[K]; J --> M[V]
        K --> N[RoPE]; L --> O[RoPE]
        N --> P["Q @ K^T / sqrt(d)"]; O --> P
        P --> Q[Masked SoftMax]; Q --> R[S @ V]; M --> R
        R --> S[Linear]; S --> T[+]; H --> T
        T --> U[RMSNorm]
        U --> V["Linear (gate)"]; U --> W["Linear (up)"]
        V --> X[SiLU]; X --> Y[×]; W --> Y
        Y --> Z["Linear (down)"]; Z --> AA[+]; T --> AA
        AA --> BB[x']
    end
 ```
 ### Autoregression
 Given a token sequence, the model predicts the probability of the next token. Each generated token is appended to the input and fed back, repeating until an end-of-sequence token or max length.
 ### Causal Mask
 ```
 sequence : [[1, 2, 3, 4, 5, 6]]
 input_ids: [[1, 2, 3, 4, 5]]
 target_ids: [[2, 3, 4, 5, 6]]
 ```
 Lower-triangular mask prevents attending to future positions:
 ```
 [[0, -inf, -inf, -inf, -inf],
 [0,    0, -inf, -inf, -inf],
 [0,    0,    0, -inf, -inf],
 [0,    0,    0,    0, -inf],
 [0,    0,    0,    0,    0]]
 ```
 ### Rotary Position Embedding (RoPE)
 RoPE embeds position into Q/K vectors via complex rotation:
 $$ q_i = R_i W_q x_i, \quad k_j = R_j W_k x_j, \quad q_i^T k_j = x_i^T W_q^T R_{i-j} W_k x_j $$
 The complex rotation `freqs_cis` is pre-computed once (`cos, sin` pairs per position). `apply_rotary_emb` multiplies Q/K as complex numbers.
 ## Training Loop
 Two-level loop: **epoch** → **batch**. Optimizer step fires every `grad_accum_steps` batches.
 ```
 on_train_begin
  on_epoch_begin
    for batch in dataloader:
      on_batch_begin
      with executor.accumulate(model):
        loss = strategy(batch)
        stand_loss = loss / executor.grad_accum_steps
        executor.backward(stand_loss)
        iteration += 1
        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_optimizer_step` | Every accumulation window | `GradientClippingCallback`, `ValidationCallback` |
 | `on_batch_end` | Every batch | `CheckpointCallback`, `MetricLoggerCallback`, `ProgressBarCallback` |
 | `on_train_end` | Training ends | `CheckpointCallback`, `MetricLoggerCallback` (final save) |
 Default callbacks (in order): `gradient_checkpointing` (activation checkpointing, optional), `checkpoint` (safetensors, rank-0), `metric_logger` (JSONL, rank-0), `progress_bar` (tqdm), `gradient_clipping`, `validation` (periodic validation on val_dataset).
 ## Strategies
 ### SEQ (Pre-training)
 Next-token cross-entropy with optional label smoothing:
 $$
 L_{\text{PT}} = -\sum_{t=1}^{T} \log P(x_t \mid x_{\lt t}; \theta)
 $$
 Keys: `input_ids`, `target_ids`
 ### SFT (Supervised Fine-Tuning)
 Masked cross-entropy (`ignore_index=-100`) over response tokens:
 $$
 L_{\text{SFT}} = -\sum_{t=P+1}^{P+L} \log P(s_t \mid s_{\lt t}; \theta)
 $$
 Keys: `input_ids`, `target_ids`, `loss_mask`
 ### DPO (Direct Preference Optimization)
 Frozen reference model, preference margin via log-ratio:
 $$
 L_{\text{DPO}} = -\mathbb{E}\left[\log\sigma\left(\beta\log\frac{\pi_\theta(y_w\mid x)}{\pi_{\text{ref}}(y_w\mid x)} - \beta\log\frac{\pi_\theta(y_l\mid x)}{\pi_{\text{ref}}(y_l\mid x)}\right)\right]
 $$
 Parameters: `beta=0.1`. Keys: `chosen`, `rejected`, `chosen_mask`, `rejected_mask`.
 ### GRPO (Group Relative Policy Optimization)
 On-policy PPO with group-normalized advantages:
 $$
 \text{Advantage}_i = \frac{r_i - \mu}{\sigma + \epsilon}
 $$
 $$
 L_{\text{GRPO}} = -\mathbb{E}\left[\min\left(\frac{\pi_\theta}{\pi_{\text{ref}}}A,\; \text{clip}\left(\frac{\pi_\theta}{\pi_{\text{ref}}}, 1-\epsilon, 1+\epsilon\right)A\right)\right] + \lambda \cdot \mathbb{E}\left[(\log\pi_\theta - \log\pi_{\text{ref}})^2\right]
 $$
 Parameters: `group_size=4`, `clip_eps=0.2`, `kl_coef=0.01`, `sync_interval=200`.
 Keys: `prompts`, `responses`, `masks`, `rewards`.
 ## LR Schedulers
 | Type | Class | Description |
 |------|-------|-------------|
 | Cosine | `CosineScheduler` | Linear warmup → cosine decay to `min_rate` |
 | SGDR | `SGDRScheduler` | Cosine annealing with warm restarts (`t_mult=2`) |
 Created by `SchedulerFactory.create(optimizer, schedule_type, **kwargs)`.
 ## Gradient Checkpointing
 Trades compute for memory by recomputing activations during backward pass. Specify module types via `gradient_checkpointing_modules`:
 ```python
 from astrai.model.components.decoder_block import DecoderBlock
 config = TrainConfig(..., gradient_checkpointing_modules=[DecoderBlock])
 ```
 Callback wraps each `DecoderBlock.forward` with `torch.utils.checkpoint.checkpoint(use_reentrant=False)`, compatible with `torch.compile`. Uses `nn.Module.apply()` for traversal — works through DDP wrappers without manual unwrap. Empty list (default) means no-op.
 ## Checkpoint
 ```
 Checkpoint(state_dict, epoch, iteration, extra, meta, config)
  ├── save(save_dir)    rank-0 only: meta.json (epoch/iteration/timestamp) + config.json (model config) + state_dict.safetensors + optional {key}.pt (optimizer.pt, scheduler.pt)
  └── load(save_dir)    broadcasts 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, ...)`
 ## Training CLI
 ```bash
 export CUDA_VISIBLE_DEVICES=0,1,2,3
 nohup python scripts/tools/train.py \
    --nprocs=4 \
    --train_type=seq \
    --data_root_path=/path/to/dataset \
    --param_path=/path/to/model \
    --batch_per_device=4 \
    --grad_accum_steps=8 \
    --warmup_ratio=0.05 \
    --max_lr=1e-4 \
    --max_grad_norm=1.0 \
    --adamw_beta1=0.9 \
    --adamw_beta2=0.95 \
    --adamw_weight_decay=0.01 \
    --window_size=2048 \
    --ckpt_interval=10000 \
    --ckpt_dir=./checkpoint \
    --random_seed=3407 \
    --label_smoothing=0.05 \
    > out.log 2> err.log &
 ```
 Full parameter reference at [params.md](params.md).
 > Document Update Time: 2026-05-28
@@ -1,34 +1,124 @@
-__version__ = "1.3.7"
+__version__ = "1.3.12"
 __author__ = "ViperEkura"
 import logging
 import os
 from astrai.config import (
    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
+from astrai.model import (
-from astrai.tokenize import AutoTokenizer
+    AutoModel,
-from astrai.trainer import CallbackFactory, SchedulerFactory, StrategyFactory, Trainer
+    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,
 )
 def setup_logging(level: str = "INFO"):
    """Attach a handler to the ``astrai`` logger (only, not root).
    Call once per process, e.g. at the top of CLI scripts.
    Set ``ASTR_LOG_LEVEL`` to override the default ``INFO``.
    """
    _logger = logging.getLogger("astrai")
    if _logger.handlers:
        return
    _level = getattr(
        logging, os.environ.get("ASTR_LOG_LEVEL", level).upper(), logging.INFO
    )
    _logger.setLevel(_level)
    _handler = logging.StreamHandler()
    _handler.setFormatter(
        logging.Formatter(
            "%(asctime)s | %(levelname)-7s | %(name)s | %(message)s",
            datefmt="%Y-%m-%d %H:%M:%S",
        )
    )
    _logger.addHandler(_handler)
 __all__ = [
    "AutoRegressiveLM",
    "AutoRegressiveLMConfig",
    "EncoderConfig",
    "TrainConfig",
    "DatasetFactory",
    "AutoTokenizer",
    "GenerationRequest",
    "InferenceEngine",
    "Trainer",
    "CallbackFactory",
    "StrategyFactory",
    "SchedulerFactory",
    "BaseFactory",
    "AutoModel",
    "AutoTokenizer",
    "BaseDataset",
    "BaseFactory",
    "BaseModelConfig",
    "BaseScheduler",
    "BaseStrategy",
    "CallbackFactory",
    "ChatTemplate",
    "Checkpoint",
    "ConfigFactory",
    "DatasetFactory",
    "EmbeddingEncoder",
    "EncoderConfig",
    "ExecutorFactory",
    "InferenceEngine",
    "LoRAConfig",
    "Pipeline",
    "PipelineConfig",
    "ProtocolHandler",
    "RDSampler",
    "SamplingPipeline",
    "SchedulerFactory",
    "Store",
    "StoreFactory",
    "StrategyFactory",
    "TrainCallback",
    "TrainConfig",
    "Trainer",
    "filter_by_length",
    "get_app",
    "get_rank",
    "get_world_size",
    "inject_lora",
    "only_on_rank",
    "run_server",
    "sample",
    "setup_logging",
    "spawn_parallel_fn",
 ]
@@ -4,13 +4,22 @@ from astrai.config.model_config import (
    ConfigFactory,
    EncoderConfig,
 )
 from astrai.config.preprocess_config import (
    InputConfig,
    OutputConfig,
    PipelineConfig,
    ProcessingConfig,
 )
 from astrai.config.train_config import TrainConfig
 __all__ = [
    # Model configuration
    "BaseModelConfig",
    "AutoRegressiveLMConfig",
    "EncoderConfig",
    "ConfigFactory",
    "TrainConfig",
    "InputConfig",
    "OutputConfig",
    "PipelineConfig",
    "ProcessingConfig",
 ]
@@ -1,86 +1,38 @@
 import json
-from dataclasses import MISSING, dataclass, fields
+from dataclasses import asdict
-from typing import Any, Dict, Optional, Self, get_type_hints
+from pathlib import Path
 from typing import Any, Dict, Self, Union
 from pydantic import ConfigDict
 from pydantic.dataclasses import dataclass
-@dataclass
+@dataclass(config=ConfigDict(use_attribute_docstrings=True))
 class BaseConfig:
    def to_dict(self) -> Dict[str, Any]:
-        d = {}
+        result = {}
-        for fld in fields(self):
+        for k, v in asdict(self).items():
-            v = getattr(self, fld.name)
+            if isinstance(v, tuple):
-            if isinstance(v, (str, int, float, bool)):
+                v = list(v)
-                d[fld.name] = v
+            try:
-            elif v is None:
+                json.dumps(v)
-                d[fld.name] = None
+                result[k] = v
-            elif isinstance(v, (dict, list, tuple)):
+            except (TypeError, ValueError):
-                try:
+                # Skip non-serializable runtime objects (e.g. model_fn, dataset).
-                    val = list(v) if isinstance(v, tuple) else v
+                # TrainConfig mixes hyperparams with callables/datasets; only the
-                    json.dumps(val)
+                # JSON-serializable subset is written to checkpoint meta.
-                    d[fld.name] = val
+                pass
-                except (TypeError, ValueError):
+        return result
                    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)
+        return cls(**d)
        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
+    @classmethod
-    def _unwrap_optional(tp) -> Optional[type]:
+    def from_file(cls, path: Union[str, Path]) -> Self:
-        if tp is None:
+        with open(path, "r", encoding="utf-8") as f:
-            return None
+            return cls.from_dict(json.load(f))
        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 to_file(self, path: Union[str, Path]):
-    def _coerce(value: Any, target_type: type) -> Any:
+        with open(path, "w", encoding="utf-8") as f:
-        if target_type is bool and isinstance(value, bool):
+            json.dump(self.to_dict(), f, indent=2, ensure_ascii=False)
            return value
        if (
            target_type is int
            and isinstance(value, (int, float))
            and not isinstance(value, bool)
        ):
            return int(value)
        if (
            target_type is float
            and isinstance(value, (int, float))
            and not isinstance(value, bool)
        ):
            return float(value)
        if target_type is str and isinstance(value, str):
            return value
        if isinstance(value, target_type):
            return value
        raise TypeError
@@ -1,10 +1,14 @@
-import json
+from typing import Any, Dict, Optional
-from dataclasses import dataclass
+
-from typing import Any, Dict, Optional, Self
+from pydantic import field_validator
 from pydantic.dataclasses import dataclass
 from astrai.config.base import BaseConfig
 from astrai.factory import BaseFactory
 _ATTN_TYPES = frozenset({"gqa", "mla"})
 _FFN_TYPES = frozenset({"mlp", "moe"})
 class ConfigFactory(BaseFactory[BaseConfig]):
    """Factory that dispatches config classes by ``model_type``."""
@@ -18,75 +22,157 @@ class ConfigFactory(BaseFactory[BaseConfig]):
@dataclass
 class BaseModelConfig(BaseConfig):
-    """Base config with ``model_type`` dispatch and file I/O."""
+    """Base config with ``model_type`` dispatch and file I/O.
    Args:
        model_type (Optional[str]): Model type identifier for AutoModel dispatch. Defaults to None.
        neftune_alpha (float): NEFTune noise alpha, 0=disabled, typical: 5.0. Defaults to 0.0.
    """
    model_type: Optional[str] = None
-
+    neftune_alpha: float = 0.0
    @classmethod
    def from_file(cls, config_path: str) -> Self:
        with open(config_path, "r") as f:
            raw: Dict[str, Any] = json.load(f)
        return cls.from_dict(raw)
    def to_file(self, config_path: str):
        d = self.to_dict()
        config_dict = {k: v for k, v in d.items() if v is not None}
        with open(config_path, "w") as f:
            json.dump(config_dict, f, indent=4)
@dataclass
@ConfigFactory.register("autoregressive_lm")
 class AutoRegressiveLMConfig(BaseModelConfig):
-    """Configuration for autoregressive language model."""
+    """Configuration for autoregressive language model.
    Args:
        model_type (Optional[str]): Model type identifier for AutoModel dispatch. Defaults to None.
        neftune_alpha (float): NEFTune noise alpha, 0=disabled, typical: 5.0. Defaults to 0.0.
        vocab_size (Optional[int]): Vocabulary size. Defaults to None.
        hidden_size (Optional[int]): Hidden dimension size. Defaults to None.
        num_hidden_layers (Optional[int]): Number of transformer layers. Defaults to None.
        rms_norm_eps (Optional[float]): Epsilon for RMSNorm. Defaults to None.
        intermediate_size (Optional[int]): Intermediate size in FFN. Defaults to None.
        tie_word_embeddings (Optional[bool]): Whether to tie embedding and lm_head weights. Defaults to None.
        max_position_embeddings (Optional[int]): Maximum sequence length the model was trained with. Defaults to None.
        rope_theta (Optional[float]): Base frequency for RoPE. Defaults to None.
        rope_scaling (Optional[dict]): RoPE scaling config, e.g. {"type": "linear", "factor": 4.0}. Defaults to None.
        attn_type (str): Attention type: 'gqa' or 'mla'. Defaults to "gqa".
        num_attention_heads (Optional[int]): Number of query attention heads. Defaults to None.
        num_key_value_heads (Optional[int]): Number of key/value heads for GQA. Defaults to None.
        use_qk_norm (Optional[bool]): Whether to apply RMSNorm to Q/K. Defaults to None.
        use_gated_attention (Optional[bool]): Whether to use gated attention. Defaults to None.
        kv_lora_rank (Optional[int]): KV compression rank, MLA only. Defaults to None.
        qk_nope_head_dim (Optional[int]): Non-RoPE head dimension, MLA only. Defaults to None.
        qk_rope_head_dim (Optional[int]): RoPE head dimension, MLA only. Defaults to None.
        ffn_type (str): FFN type: 'mlp' or 'moe'. Defaults to "mlp".
        n_routed_experts (Optional[int]): Number of routed experts, MoE only. Defaults to None.
        n_shared_experts (Optional[int]): Number of shared experts, MoE only. Defaults to None.
        n_activated_experts (Optional[int]): Number of activated experts per token, MoE only. Defaults to None.
        topk_method (Optional[str]): Top-k routing method, MoE only. Defaults to None.
        moe_intermediate_size (Optional[int]): Expert hidden dim, defaults to intermediate_size if None. MoE only.
        shared_expert_intermediate_size (Optional[int]): Shared expert hidden dim, defaults to intermediate_size if None. MoE only.
        norm_topk_prob (bool): Normalize top-k routing probabilities. Defaults to True.
        decoder_sparse_step (int): Frequency of MoE layers, 1=every layer. Defaults to 1.
        mlp_only_layers (Optional[list[int]]): Layer indices using dense MLP instead of MoE. Defaults to None.
    """
    vocab_size: Optional[int] = None
-    dim: Optional[int] = None
+    hidden_size: Optional[int] = None
-    n_layers: Optional[int] = None
+    num_hidden_layers: Optional[int] = None
-    norm_eps: Optional[float] = None
+    rms_norm_eps: Optional[float] = None
-    dim_ffn: Optional[int] = None
+    intermediate_size: Optional[int] = None
-    tie_weight: Optional[bool] = None
+    tie_word_embeddings: Optional[bool] = None
-
+    max_position_embeddings: 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
+    num_attention_heads: Optional[int] = None
-    n_kv_heads: Optional[int] = None
+    num_key_value_heads: Optional[int] = None
    use_qk_norm: Optional[bool] = None
    use_gated_attention: Optional[bool] = None
    kv_lora_rank: Optional[int] = None
    qk_nope_head_dim: Optional[int] = None
    qk_rope_head_dim: Optional[int] = None
    ffn_type: str = "mlp"
    n_routed_experts: Optional[int] = None
    n_shared_experts: Optional[int] = None
    n_activated_experts: Optional[int] = None
    topk_method: Optional[str] = None
    moe_intermediate_size: Optional[int] = None
    shared_expert_intermediate_size: Optional[int] = None
    norm_topk_prob: bool = True
    decoder_sparse_step: int = 1
    mlp_only_layers: Optional[list[int]] = None
    moe_aux_loss_coef: float = 0.01
    @field_validator("attn_type")
    def _validate_attn_type(cls, v: str) -> str:
        if v not in _ATTN_TYPES:
            raise ValueError(
                f"attn_type must be one of {sorted(_ATTN_TYPES)}, got {v!r}"
            )
        return v
    @field_validator("ffn_type")
    def _validate_ffn_type(cls, v: str) -> str:
        if v not in _FFN_TYPES:
            raise ValueError(f"ffn_type must be one of {sorted(_FFN_TYPES)}, got {v!r}")
        return v
    @field_validator("decoder_sparse_step")
    def _validate_decoder_sparse_step(cls, v: int) -> int:
        if v < 1:
            raise ValueError(f"decoder_sparse_step must be at least 1, got {v}")
        return v
@dataclass
@ConfigFactory.register("embedding")
 class EncoderConfig(BaseModelConfig):
-    """Configuration for embedding encoder model."""
+    """Configuration for embedding encoder model.
    Args:
        model_type (Optional[str]): Model type identifier for AutoModel dispatch. Defaults to None.
        neftune_alpha (float): NEFTune noise alpha, 0=disabled, typical: 5.0. Defaults to 0.0.
        vocab_size (Optional[int]): Vocabulary size. Defaults to None.
        hidden_size (Optional[int]): Hidden dimension size. Defaults to None.
        num_hidden_layers (Optional[int]): Number of transformer layers. Defaults to None.
        rms_norm_eps (Optional[float]): Epsilon for RMSNorm. Defaults to None.
        intermediate_size (Optional[int]): Intermediate size in FFN. Defaults to None.
        max_position_embeddings (Optional[int]): Maximum sequence length the model was trained with. Defaults to None.
        rope_theta (Optional[float]): Base frequency for RoPE. Defaults to None.
        rope_scaling (Optional[dict]): RoPE scaling config, e.g. {"type": "linear", "factor": 4.0}. Defaults to None.
        attn_type (str): Attention type: 'gqa' or 'mla'. Defaults to "gqa".
        num_attention_heads (Optional[int]): Number of query attention heads. Defaults to None.
        num_key_value_heads (Optional[int]): Number of key/value heads for GQA. Defaults to None.
        use_qk_norm (Optional[bool]): Whether to apply RMSNorm to Q/K. Defaults to None.
        use_gated_attention (Optional[bool]): Whether to use gated attention. Defaults to None.
        ffn_type (str): FFN type: 'mlp' or 'moe'. Defaults to "mlp".
        pooling_type (Optional[str]): Pooling strategy for embedding, e.g. 'mean', 'cls'. Defaults to None.
        normalize_embeddings (Optional[bool]): Whether to L2-normalize output embeddings. Defaults to None.
    """
    vocab_size: Optional[int] = None
-    dim: Optional[int] = None
+    hidden_size: Optional[int] = None
-    n_layers: Optional[int] = None
+    num_hidden_layers: Optional[int] = None
-    norm_eps: Optional[float] = None
+    rms_norm_eps: Optional[float] = None
-    dim_ffn: Optional[int] = None
+    intermediate_size: Optional[int] = None
-
+    max_position_embeddings: 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
+    num_attention_heads: Optional[int] = None
-    n_kv_heads: Optional[int] = None
+    num_key_value_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
    @field_validator("attn_type")
    def _validate_attn_type(cls, v: str) -> str:
        if v not in _ATTN_TYPES:
            raise ValueError(
                f"attn_type must be one of {sorted(_ATTN_TYPES)}, got {v!r}"
            )
        return v
    @field_validator("ffn_type")
    def _validate_ffn_type(cls, v: str) -> str:
        if v not in _FFN_TYPES:
            raise ValueError(f"ffn_type must be one of {sorted(_FFN_TYPES)}, got {v!r}")
        return v
@@ -0,0 +1,152 @@
 """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 field
 from typing import Dict, List, Optional
 from pydantic import field_validator
 from pydantic.dataclasses import dataclass
 from astrai.config.base import BaseConfig
 _PACKING_STRATEGIES = frozenset({"simple", "bfd", "bfd_split"})
 _TRUNCATION_MODES = frozenset({"keep_start", "keep_end"})
 _STORAGE_FORMATS = frozenset({"bin", "jsonl"})
 _POSITION_IDS_MODES = frozenset({"none", "doc_reset", "continuous"})
@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", ...}]},
        }}}
    Args:
        sections (Optional[List[Dict]]): Section list for single-output mode. Defaults to None.
        sources (Optional[Dict[str, Dict]]): Source map for multi-output mode, DPO/GRPO. Defaults to None.
    """
    sections: Optional[List[Dict]] = None
    sources: Optional[Dict[str, Dict]] = None
@dataclass
 class ProcessingConfig(BaseConfig):
    """Processing configuration for tokenization and packing.
    Args:
        max_seq_len (int): Maximum sequence length. Defaults to 2048.
        min_chars (int): Minimum number of characters to keep. Defaults to 50.
        max_chars (int): Maximum number of characters to keep. Defaults to 2_000_000.
        max_items (Optional[int]): Maximum number of items to process, None=unlimited. Defaults to None.
        batch_size (int): Number of records tokenized together. Defaults to 256.
        packing_strategy (str): How to pack sequences: 'simple', 'bfd', or 'bfd_split'. Defaults to "simple".
        max_packed_len (int): Maximum length of a packed bin. Defaults to 8192.
        truncation_mode (str): How to truncate over-length sequences: 'keep_start' or 'keep_end'. Defaults to "keep_start".
    """
    max_seq_len: int = 2048
    min_chars: int = 50
    max_chars: int = 2_000_000
    max_items: Optional[int] = None
    batch_size: int = 256
    packing_strategy: str = "simple"
    max_packed_len: int = 8192
    truncation_mode: str = "keep_start"
    @field_validator("packing_strategy")
    def _validate_packing_strategy(cls, v: str) -> str:
        if v not in _PACKING_STRATEGIES:
            raise ValueError(
                f"packing_strategy must be one of {sorted(_PACKING_STRATEGIES)}, got {v!r}"
            )
        return v
    @field_validator("truncation_mode")
    def _validate_truncation_mode(cls, v: str) -> str:
        if v not in _TRUNCATION_MODES:
            raise ValueError(
                f"truncation_mode must be one of {sorted(_TRUNCATION_MODES)}, got {v!r}"
            )
        return v
    @field_validator("max_seq_len", "batch_size", "max_packed_len")
    def _validate_positive_int(cls, v: int) -> int:
        if v <= 0:
            raise ValueError(f"must be positive, got {v}")
        return v
    @field_validator("min_chars")
    def _validate_non_negative(cls, v: int) -> int:
        if v < 0:
            raise ValueError(f"min_chars must be non-negative, got {v}")
        return v
@dataclass
 class OutputConfig(BaseConfig):
    """Output configuration for storage.
    Args:
        domain_key (Optional[str]): Domain key for the output store. Defaults to None.
        storage_format (str): Storage format: 'bin' or 'jsonl'. Defaults to "bin".
        max_tokens_per_shard (int): Maximum tokens per shard before splitting. Defaults to 100_000_000.
        dtype (Dict[str, str]): Per-key dtype overrides, e.g. {"input_ids": "int32"}. Defaults to {}.
        position_ids_mode (str): Position ids mode: 'none', 'doc_reset', or 'continuous'. Defaults to "doc_reset".
    """
    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"
    @field_validator("storage_format")
    def _validate_storage_format(cls, v: str) -> str:
        if v not in _STORAGE_FORMATS:
            raise ValueError(
                f"storage_format must be one of {sorted(_STORAGE_FORMATS)}, got {v!r}"
            )
        return v
    @field_validator("position_ids_mode")
    def _validate_position_ids_mode(cls, v: str) -> str:
        if v not in _POSITION_IDS_MODES:
            raise ValueError(
                f"position_ids_mode must be one of {sorted(_POSITION_IDS_MODES)}, got {v!r}"
            )
        return v
@dataclass
 class PipelineConfig(BaseConfig):
    """Top-level preprocessing pipeline config.
    Args:
        version (int): Config schema version. Defaults to 1.
        input (InputConfig): Input mapping config.
        mask (Dict[str, str]): Per-field mask labels, e.g. {"system": "mask", "assistant": "train"}. Defaults to {}.
        mask_default (str): Default mask label for unlisted fields. Defaults to "mask".
        preprocessing (ProcessingConfig): Processing config.
        output (OutputConfig): Output config.
    """
    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,7 +1,9 @@
-from dataclasses import dataclass, field, fields
+from dataclasses import field
-from typing import Callable, List, Optional
+from typing import Any, Callable, Dict, List, Optional
 import torch.nn as nn
 from pydantic import ConfigDict, field_validator, model_validator
 from pydantic.dataclasses import dataclass
 from torch.optim import Optimizer
 from torch.optim.lr_scheduler import LRScheduler
 from torch.utils.data import Dataset
@@ -9,132 +11,208 @@ from torch.utils.data import Dataset
 from astrai.config.base import BaseConfig
 from astrai.model.components.lora import LoRAConfig
-
+_TRAIN_TYPES = frozenset({"seq", "sft", "dpo", "grpo", "online_grpo", "online_dpo"})
-def required(**kw):
+_PARALLEL_MODES = frozenset({"none", "ddp", "fsdp"})
-    return {"required": True, **kw}
+_BACKENDS = frozenset({"nccl", "gloo"})
 _START_METHODS = frozenset({"spawn", "fork", "forkserver"})
 _COMPILE_MODES = frozenset({"default", "reduce-overhead", "max-autotune"})
-@dataclass
+@dataclass(config=ConfigDict(arbitrary_types_allowed=True))
 class TrainConfig(BaseConfig):
-    # basic setting
+    """Training configuration.
    model_fn: Callable[[], nn.Module] = field(
        default=None, metadata=required(help="Model factory for training.")
    )
    strategy: str = field(default=None, metadata=required(help="Training strategy."))
    dataset: Dataset = field(
        default=None, metadata=required(help="Dataset for training.")
    )
    optimizer_fn: Callable[[nn.Module], Optimizer] = field(
        default=None, metadata=required(help="Optimizer factory for training.")
    )
    scheduler_fn: Callable[[Optimizer], LRScheduler] = field(
        default=None, metadata=required(help="Scheduler factory for training.")
    )
    n_epoch: int = field(default=1, metadata={"help": "Number of epochs for training."})
    batch_per_device: int = field(
        default=4, metadata={"help": "Batch size per device."}
    )
    grad_accum_steps: int = field(
        default=1, metadata={"help": "Number of iterations between steps."}
    )
    max_grad_norm: float = field(
        default=1.0, metadata={"help": "Maximum gradient norm."}
    )
    gradient_checkpointing_modules: list = field(
        default_factory=list,
        metadata={"help": "Module types to enable activation checkpointing for."},
    )
-    # checkpoint setting
+    Combines hyperparameters with runtime objects (model_fn, dataset, etc.).
-    start_epoch: int = field(default=0, metadata={"help": "Start epoch for training."})
+    Only JSON-serializable fields are written to checkpoint meta via to_dict().
    start_batch: int = field(
        default=0, metadata={"help": "Start batch iteration for training."}
    )
    ckpt_dir: str = field(
        default="./checkpoint", metadata={"help": "Checkpoint directory."}
    )
    ckpt_interval: int = field(
        default=5000, metadata={"help": "Number of iterations between checkpoints."}
    )
-    # lora setting
+    Args:
-    lora: Optional[LoRAConfig] = field(
+        model_fn (Callable[[], nn.Module]): Model factory for training.
-        default=None,
+        strategy (str): Training strategy (seq, sft, dpo, grpo, online_*).
-        metadata={"help": "LoRA config. None means full fine-tuning."},
+        dataset (Dataset): Dataset for training.
-    )
+        optimizer_fn (Callable[[nn.Module], Optimizer]): Optimizer factory for training.
        optimizer_name (Optional[str]): Serializable built-in optimizer identifier. Defaults to None.
        optimizer_hyperparameters (Dict[str, Any]): Serializable optimizer settings. Defaults to {}.
        scheduler_fn (Callable[[Optimizer], LRScheduler]): Scheduler factory for training.
        n_epoch (int): Number of epochs for training. Defaults to 1.
        batch_per_device (int): Batch size per device. Defaults to 4.
        grad_accum_steps (int): Number of iterations between optimizer steps. Defaults to 1.
        max_grad_norm (Optional[float]): Maximum gradient norm. None disables clipping. Defaults to 1.0.
        gradient_checkpointing_modules (List[type]): Module types to enable activation checkpointing for. Defaults to [].
        compile_mode (Optional[str]): torch.compile mode: 'default', 'reduce-overhead', 'max-autotune', or None. Defaults to None.
        start_epoch (int): Start epoch for training. Defaults to 0.
        start_samples (int): Start samples count (per rank). Superseded by checkpoint consumed_samples. Defaults to 0.
        ckpt_dir (str): Checkpoint directory. Defaults to "./checkpoint".
        ckpt_interval (int): Number of optimizer steps between checkpoints. Defaults to 5000.
        lora (Optional[LoRAConfig]): LoRA config. None means full fine-tuning. Defaults to None.
        metrics (List[str]): Metrics to record during training. Defaults to ["loss", "lr", "grad_norm"].
        random_seed (int): Random seed. Defaults to 3407.
        num_workers (int): Number of workers for dataloader. Defaults to 0.
        prefetch_factor (Optional[int]): Prefetch factor for dataloader. Defaults to None.
        pin_memory (bool): Pin memory for dataloader. Defaults to False.
        collate_fn (Optional[Callable[[List[Any]], Any]]): Collate function for dataloader (e.g. dpo_collate_fn). Defaults to None.
        nprocs (int): Number of processes for distributed training. Defaults to 1.
        backend (str): Distributed training backend. Defaults to "nccl".
        master_addr (str): Master address for distributed training. Defaults to "localhost".
        master_port (str): Master port for distributed training. Defaults to "29500".
        parallel_mode (str): Parallel strategy: none, ddp, fsdp. Defaults to "none".
        start_method (str): Multiprocessing start method: spawn/fork/forkserver. Defaults to "spawn".
        device_type (str): Device type for distributed training. Defaults to "cuda".
        val_dataset (Optional[Dataset]): Dataset for validation. Defaults to None.
        val_split (Optional[float]): Ratio to split from training dataset for validation, e.g. 0.05. Defaults to None.
        val_step (int): Number of optimizer steps between validation runs. Defaults to 1000.
        neftune_alpha (float): NEFTune noise alpha, 0=disabled, typical: 5.0. Defaults to 0.0.
        moe_aux_loss_coef (float): Weight applied to the MoE load-balancing loss. Defaults to 0.01.
        rollout_interval (int): Number of optimizer steps between online rollouts. Defaults to 512.
        rollout_temperature (float): Sampling temperature for online rollout. Defaults to 0.7.
        rollout_top_k (int): Top-k filtering for online rollout, 0=disable. Defaults to 0.
        rollout_top_p (float): Top-p (nucleus) filtering for online rollout. Defaults to 0.9.
        rollout_max_tokens (int): Maximum generated tokens per response in rollout. Defaults to 1024.
        reward_model_fn (Optional[Callable]): Factory for reward model, required for online RL strategies. Defaults to None.
        executor_kwargs (Dict[str, Any]): Extra kwargs passed to ExecutorFactory.create(). Defaults to {}.
        extra_kwargs (Dict[str, Any]): Other arguments. Defaults to {}.
    """
-    # metric setting
+    model_fn: Callable[[], nn.Module]
-    log_dir: str = field(
+    strategy: str
-        default="./checkpoint/logs", metadata={"help": "Directory for metric logs."}
+    dataset: Dataset
-    )
+    optimizer_fn: Callable[[nn.Module], Optimizer]
-    log_interval: int = field(
+    scheduler_fn: Callable[[Optimizer], LRScheduler]
-        default=100,
+    optimizer_name: Optional[str] = None
-        metadata={"help": "Number of batch iterations between metric logs."},
+    optimizer_hyperparameters: Dict[str, Any] = field(default_factory=dict)
-    )
+    n_epoch: int = 1
-    metrics: List[str] = field(
+    batch_per_device: int = 4
-        default_factory=lambda: ["loss", "lr"],
+    grad_accum_steps: int = 1
-        metadata={"help": "Metrics to record during training."},
+    max_grad_norm: Optional[float] = 1.0
-    )
+    gradient_checkpointing_modules: List[type] = field(default_factory=list)
    compile_mode: Optional[str] = None
-    # dataloader setting
+    start_epoch: int = 0
-    random_seed: int = field(default=3407, metadata={"help": "Random seed."})
+    start_samples: int = 0
-    num_workers: int = field(
+    ckpt_dir: str = "./checkpoint"
-        default=0, metadata={"help": "Number of workers for dataloader."}
+    ckpt_interval: int = 5000
    )
    prefetch_factor: Optional[int] = field(
        default=None, metadata={"help": "Prefetch factor for dataloader."}
    )
    pin_memory: bool = field(
        default=False, metadata={"help": "Pin memory for dataloader."}
    )
-    # distributed training
+    lora: Optional[LoRAConfig] = None
    nprocs: int = field(
        default=1, metadata={"help": "Number of processes for distributed training."}
    )
    backend: str = field(
        default="nccl", metadata={"help": "Distributed training backend."}
    )
    master_addr: str = field(
        default="localhost",
        metadata={"help": "Master address for distributed training."},
    )
    master_port: str = field(
        default="29500", metadata={"help": "Master port for distributed training."}
    )
    parallel_mode: str = field(
        default="none",
        metadata={"help": "Parallel strategy: none, ddp, fsdp."},
    )
    start_method: str = field(
        default="spawn",
        metadata={"help": "Multiprocessing start method (spawn/fork/forkserver)."},
    )
-    # others
+    metrics: List[str] = field(default_factory=lambda: ["loss", "lr", "grad_norm"])
    device_type: str = field(
        default="cuda", metadata={"help": "Device type for distributed training."}
    )
    val_dataset: Optional[Dataset] = field(
        default=None, metadata={"help": "Dataset for validation."}
    )
    val_step: int = field(
        default=1000,
        metadata={"help": "Number of optimizer steps between validation runs."},
    )
-    executor_kwargs: dict = field(
+    random_seed: int = 3407
-        default_factory=dict,
+    num_workers: int = 0
-        metadata={"help": "Extra kwargs passed to ExecutorFactory.create()."},
+    prefetch_factor: Optional[int] = None
-    )
+    pin_memory: bool = False
-    extra_kwargs: dict = field(
+    collate_fn: Optional[Callable[[List[Any]], Any]] = None
        default_factory=dict, metadata={"help": "Other arguments."}
    )
-    def __post_init__(self):
+    nprocs: int = 1
-        self.validate()
+    backend: str = "nccl"
    master_addr: str = "localhost"
    master_port: str = "29500"
    parallel_mode: str = "none"
    start_method: str = "spawn"
-    def validate(self):
+    device_type: str = "cuda"
-        for fld in fields(self):
+    val_dataset: Optional[Dataset] = None
-            if fld.metadata.get("required") and getattr(self, fld.name) is None:
+    val_split: Optional[float] = None
-                raise ValueError(f"TrainConfig.{fld.name} is required but got None.")
+    val_step: int = 1000
    neftune_alpha: float = 0.0
    moe_aux_loss_coef: float = 0.01
    rollout_interval: int = 512
    rollout_temperature: float = 0.7
    rollout_top_k: int = 0
    rollout_top_p: float = 0.9
    rollout_max_tokens: int = 1024
    reward_model_fn: Optional[Callable] = None
    executor_kwargs: Dict[str, Any] = field(default_factory=dict)
    extra_kwargs: Dict[str, Any] = field(default_factory=dict)
    @field_validator("strategy")
    def _validate_strategy(cls, v: str) -> str:
        if v not in _TRAIN_TYPES:
            raise ValueError(
                f"strategy must be one of {sorted(_TRAIN_TYPES)}, got {v!r}"
            )
        return v
    @field_validator("parallel_mode")
    def _validate_parallel_mode(cls, v: str) -> str:
        if v not in _PARALLEL_MODES:
            raise ValueError(
                f"parallel_mode must be one of {sorted(_PARALLEL_MODES)}, got {v!r}"
            )
        return v
    @field_validator("backend")
    def _validate_backend(cls, v: str) -> str:
        if v not in _BACKENDS:
            raise ValueError(f"backend must be one of {sorted(_BACKENDS)}, got {v!r}")
        return v
    @field_validator("start_method")
    def _validate_start_method(cls, v: str) -> str:
        if v not in _START_METHODS:
            raise ValueError(
                f"start_method must be one of {sorted(_START_METHODS)}, got {v!r}"
            )
        return v
    @field_validator("compile_mode")
    def _validate_compile_mode(cls, v: Optional[str]) -> Optional[str]:
        if v is not None and v not in _COMPILE_MODES:
            raise ValueError(
                f"compile_mode must be one of {sorted(_COMPILE_MODES)} or None, got {v!r}"
            )
        return v
    @field_validator(
        "n_epoch",
        "batch_per_device",
        "grad_accum_steps",
        "ckpt_interval",
        "val_step",
        "rollout_interval",
        "rollout_max_tokens",
    )
    def _validate_positive_int(cls, v: int) -> int:
        if v <= 0:
            raise ValueError(f"must be positive, got {v}")
        return v
    @field_validator("rollout_temperature")
    def _validate_positive_float(cls, v: float) -> float:
        if v <= 0:
            raise ValueError(f"must be positive, got {v}")
        return v
    @field_validator("rollout_top_p")
    def _validate_top_p(cls, v: float) -> float:
        if not 0 < v <= 1:
            raise ValueError(f"rollout_top_p must be in (0, 1], got {v}")
        return v
    @field_validator(
        "rollout_top_k", "num_workers", "neftune_alpha", "moe_aux_loss_coef"
    )
    def _validate_non_negative(cls, v):
        if v < 0:
            raise ValueError(f"must be non-negative, got {v}")
        return v
    @field_validator("max_grad_norm")
    def _validate_max_grad_norm(cls, v: Optional[float]) -> Optional[float]:
        if v is not None and v <= 0:
            raise ValueError(f"max_grad_norm must be positive or None, got {v}")
        return v
    @field_validator("val_split")
    def _validate_val_split(cls, v: Optional[float]) -> Optional[float]:
        if v is not None and not 0 < v < 1:
            raise ValueError(f"val_split must be in (0, 1) or None, got {v}")
        return v
    @model_validator(mode="after")
    def _validate_online_strategy(self) -> "TrainConfig":
        if self.strategy.startswith("online_") and self.reward_model_fn is None:
            raise ValueError(
                f"reward_model_fn is required for online RL strategy {self.strategy!r}"
            )
        return self
@@ -1,31 +1,37 @@
 from astrai.dataset.dataset import (
    BaseDataset,
    DatasetFactory,
    dpo_collate_fn,
    grpo_collate_fn,
 )
-from astrai.dataset.sampler import ResumableDistributedSampler
+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,
 )
 __all__ = [
    "BaseDataset",
    "DatasetFactory",
    "dpo_collate_fn",
    "grpo_collate_fn",
    "Store",
    "Streamable",
    "Recordable",
    "StoreFactory",
    "H5Store",
    "MmapStore",
    "JsonlStore",
    "detect_format",
    "save_h5",
    "load_h5",
    "save_bin",
    "load_bin",
-    "ResumableDistributedSampler",
+    "RDSampler",
 ]
@@ -1,7 +1,31 @@
-"""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.
 """
 from abc import ABC, abstractmethod
-from typing import Dict, List, Optional
+from functools import partial
 from typing import Callable, Dict, List, Optional
 import torch
 from torch import Tensor
@@ -13,296 +37,470 @@ from astrai.dataset.storage import (
    detect_format,
 )
 from astrai.factory import BaseFactory
 from astrai.tokenize import AutoTokenizer
 def dpo_tokenize(
    record: dict,
    tokenizer,
    max_len: int = 2048,
 ) -> Optional[dict]:
    """Tokenize one DPO record into chosen/rejected + masks.
    Applies the tokenizer's chat template so token sequences match the
    SFT checkpoint's format.  Prompt is rendered with
    ``add_generation_prompt=True``; chosen/rejected are appended as a
    single assistant turn.
    Accepts:
    - Flat:   ``{"prompt": str, "chosen": str, "rejected": str}``
    - Conv:   ``{"prompt": [{role, content}, ...], "chosen": [...], ...}``
    - Legacy: ``{"input": str, "chosen": str, "rejected": str}``
    No packing, no ``position_ids`` — DPO sequences are independent.
    """
    prompt = record.get("prompt") or record.get("input")
    chosen = record.get("chosen")
    rejected = record.get("rejected")
    if prompt is None or chosen is None or rejected is None:
        return None
    prompt_messages = _to_messages(prompt)
    chosen_text = _extract_text(chosen)
    rejected_text = _extract_text(rejected)
    if chosen_text is None or rejected_text is None:
        return None
    chosen_messages = prompt_messages + [{"role": "assistant", "content": chosen_text}]
    rejected_messages = prompt_messages + [
        {"role": "assistant", "content": rejected_text}
    ]
    prompt_ids = tokenizer.apply_chat_template(
        prompt_messages, tokenize=True, add_generation_prompt=True
    )
    ch_ids = tokenizer.apply_chat_template(
        chosen_messages, tokenize=True, add_generation_prompt=False
    )
    re_ids = tokenizer.apply_chat_template(
        rejected_messages, tokenize=True, add_generation_prompt=False
    )
    full_ch = ch_ids[:max_len]
    full_re = re_ids[:max_len]
    prompt_len = min(len(prompt_ids), max_len)
    ch_mask = [0] * prompt_len + [1] * max(0, len(full_ch) - prompt_len)
    ch_mask = ch_mask[:max_len]
    re_mask = [0] * prompt_len + [1] * max(0, len(full_re) - prompt_len)
    re_mask = re_mask[:max_len]
    return {
        "chosen": full_ch,
        "rejected": full_re,
        "chosen_mask": ch_mask,
        "rejected_mask": re_mask,
    }
 def _to_messages(value) -> list:
    """Accept str or conversation list; return message list."""
    if isinstance(value, str):
        return [{"role": "user", "content": value}]
    if isinstance(value, list):
        return value
    return [{"role": "user", "content": str(value)}]
 def _extract_text(value) -> Optional[str]:
    """Accept str or conversation list; return plain text."""
    if value is None:
        return None
    if isinstance(value, str):
        return value
    if isinstance(value, list):
        return "".join(m.get("content", "") for m in value if isinstance(m, dict))
    return None
 def dpo_processor(
    record: dict,
    tokenizer,
    max_len: int = 2048,
 ) -> Dict[str, Tensor]:
    """DPO processor: wraps :func:`dpo_tokenize` and returns tensors."""
    result = dpo_tokenize(record, tokenizer, max_len=max_len)
    if result is None:
        raise ValueError(f"Malformed DPO record: {list(record.keys())}")
    return {
        "chosen": torch.tensor(result["chosen"], dtype=torch.int32),
        "rejected": torch.tensor(result["rejected"], dtype=torch.int32),
        "chosen_mask": torch.tensor(result["chosen_mask"], dtype=torch.bool),
        "rejected_mask": torch.tensor(result["rejected_mask"], dtype=torch.bool),
    }
 def dpo_collate_fn(batch: List[Dict[str, Tensor]]) -> Dict[str, Tensor]:
    """Collate variable-length DPO samples into padded 2-D tensors.
    Input: list of dicts, each with:
      - chosen:        [C_i]
      - rejected:      [R_i]
      - chosen_mask:   [C_i]
      - rejected_mask: [R_i]
    Output (padded to the max length across chosen/rejected within the batch):
      - chosen:        [B, S_max]
      - rejected:      [B, S_max]
      - chosen_mask:   [B, S_max]
      - rejected_mask: [B, S_max]
    """
    B = len(batch)
    S_max = max(b["chosen"].size(0) for b in batch)
    S_max = max(S_max, max(b["rejected"].size(0) for b in batch))
    chosen = torch.zeros(B, S_max, dtype=torch.long)
    rejected = torch.zeros(B, S_max, dtype=torch.long)
    chosen_mask = torch.zeros(B, S_max, dtype=torch.bool)
    rejected_mask = torch.zeros(B, S_max, dtype=torch.bool)
    for i, b in enumerate(batch):
        c_len = b["chosen"].size(0)
        r_len = b["rejected"].size(0)
        chosen[i, :c_len] = b["chosen"]
        rejected[i, :r_len] = b["rejected"]
        chosen_mask[i, :c_len] = b["chosen_mask"]
        rejected_mask[i, :r_len] = b["rejected_mask"]
    return {
        "chosen": chosen,
        "rejected": rejected,
        "chosen_mask": chosen_mask,
        "rejected_mask": rejected_mask,
    }
 def grpo_collate_fn(batch: List[Dict[str, Tensor]]) -> Dict[str, Tensor]:
    """Collate variable-length GRPO samples into padded 3-D tensors.
    Input: list of dicts, each with:
      - prompts:  [P_i]
      - responses: list of G tensors, each [R_ij]
      - masks:     list of G tensors, each [R_ij]
      - rewards:  [G]
    Output:
      - prompts:   [B, P_max], left-padded
      - prompt_mask: [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"])
    prompts = torch.zeros(B, P_max, dtype=torch.long)
    prompt_mask = torch.zeros(B, P_max, dtype=torch.bool)
    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)
    for i, b in enumerate(batch):
        p_len = b["prompts"].size(0)
        prompts[i, -p_len:] = b["prompts"]
        prompt_mask[i, -p_len:] = True
        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,
        "prompt_mask": prompt_mask,
        "responses": responses,
        "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 all dataset types.
+    """Abstract base class for dataset types.
-    Implements common functionality for window-based data fetching.
+    Holds a :class:`Store`.  All sample-id indexing is delegated to the
-    Uses a storage abstraction for format-agnostic data loading.
+    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).
    """
-    def __init__(self, window_size: int, stride: int):
+    required_keys: List[str] = []
    def __init__(self, store: Store):
        super().__init__()
-        self.window_size = window_size
+        self.store: Store = store
-        self.stride = stride
+        validate_keys(store, self.required_keys)
        self.storage: Optional[Store] = None
-    @property
+    def __len__(self) -> int:
-    def required_keys(self) -> List[str]:
+        return len(self.store)
        """Return required storage keys for this dataset type.
        Subclasses should override to specify expected keys.
        """
        return []
    def _validate_keys(self):
        if not self.required_keys:
            return
        actual_keys = set(self.storage.keys)
        missing = [k for k in self.required_keys if k not in actual_keys]
        if missing:
            raise KeyError(
                f"Dataset {type(self).__name__} requires keys {self.required_keys}, "
                f"but storage at {self._load_path} only has {sorted(actual_keys)}. "
                f"Missing: {missing}"
            )
    def load(self, load_path: str, storage_type: Optional[str] = None):
        """Load dataset from the given path.
        Auto-detects the storage format if not specified.
        Args:
            load_path: Path to the data directory or file
            storage_type: Force a specific storage type ("h5", "bin"),
                          or None for auto-detection
        Raises:
            KeyError: If the loaded storage is missing required keys.
        """
        if storage_type is None:
            storage_type = detect_format(load_path)
        self.storage = StoreFactory.create(storage_type)
        self._load_path = load_path
        self.storage.load(load_path)
        self._validate_keys()
    @property
    def count(self) -> int:
        """Return the total number of raw elements (tokens) in the dataset."""
        if self.storage is None:
            return 0
        return len(self.storage)
    @property
    def keys(self) -> List[str]:
-        """Return the available data keys."""
+        return self.store.keys
        if self.storage is None:
            return []
        return self.storage.keys
-    def get_index(self, index: int) -> tuple:
+    @property
-        """Calculate begin and end indices for a sample.
+    def token_count(self) -> int:
-
+        return self.store.token_count
        Args:
            index: Sample index
        Returns:
            Tuple of (begin_idx, end_idx)
        """
        if self.storage is None:
            raise RuntimeError("Dataset not loaded, call load() first")
        total = len(self.storage)
        if total <= self.window_size:
            raise ValueError(
                f"Data too short: {total} tokens <= window_size {self.window_size}"
            )
        begin_idx = min(index * self.stride, total - 1 - self.window_size)
        end_idx = min(begin_idx + self.window_size, total - 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:
        if self.storage is None:
            return 0
        total = len(self.storage)
        if total <= self.window_size:
            return 0
        return (total - 1 - self.window_size) // self.stride + 1
 class DatasetFactory(BaseFactory["BaseDataset"]):
-    """Factory class for creating dataset instances.
+    """Factory for creating dataset instances by train-type.
-    Supports decorator-based registration for extensible dataset types.
+    Use :meth:`DatasetFactory.register("custom")` to register new
-    All default dataset types (seq, sft, dpo, grpo) are registered automatically
+    dataset classes; they must inherit from :class:`BaseDataset`.
    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):
        """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,
+        load_path: Optional[str] = None,
-        window_size: int,
+        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: Type of training dataset
+            train_type: Registered dataset name ("seq", "sft", "dpo",
-            load_path: Path to the data file
+                "grpo", …).
-            window_size: Window size for data sampling
+            load_path: Path to the data file or directory (ignored if
-            stride: Stride between consecutive samples (default: same as window_size)
+                *store* is given).
-            storage_type: Storage type ("h5", "bin") or None for auto-detection
+            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 ("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
+            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
-        dataset = cls.create(train_type, window_size, stride)
+        processor = cls._maybe_build_processor(
-        dataset.load(load_path, storage_type=storage_type)
+            train_type, storage_type, tokenizer_path, max_len
        )
-        return dataset
+        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:
            load_kwargs = dict(kwargs)
            if (
                tokenizer_path is not None
                and storage_type == "jsonl"
                and train_type in ("seq", "sft")
                and "tokenizer_path" not in load_kwargs
            ):
                load_kwargs["tokenizer_path"] = tokenizer_path
            store.load(load_path, **load_kwargs)
-    @classmethod
+        return cls.create(train_type, store=store)
-    def available_types(cls) -> list:
+
-        """Return list of registered dataset type names."""
+    @staticmethod
-        return cls.list_registered()
+    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 (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."""
+    """Dataset for sequential next-token prediction training.
-    def __init__(self, window_size: int, stride: int):
+    Stream mode: ``store.fetch(begin, end, "sequence")`` returns the
-        super().__init__(window_size, stride)
+    input window; the +1 shifted call returns the next-token target.
    """
-    @property
+    required_keys = ["sequence"]
    def required_keys(self) -> List[str]:
        return ["sequence"]
-    def _fetch_data(self, begin_idx: int, end_idx: int) -> Tensor:
+    def __getitem__(self, index: int):
-        return self.storage.fetch(begin_idx, end_idx, "sequence")
+        begin, end = self.store.sample_window(index)
-
+        x = self.store.fetch(begin, end, "sequence")
-    def __getitem__(self, index):
+        y = self.store.fetch(begin + 1, end + 1, "sequence")
-        begin_idx, end_idx = self.get_index(index)
+        return {
-
+            "input_ids": x.to(dtype=torch.long),
-        x = self._fetch_data(begin_idx, end_idx).to(dtype=torch.long)
+            "target_ids": y.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."""
+    """Dataset for supervised fine-tuning with loss masking.
-    def __init__(self, window_size: int, stride: int):
+    Stream mode: ``sequence``/``loss_mask``/``position_ids`` are sliced
-        super().__init__(window_size, stride)
+    to the window.  ``loss_mask`` and ``target_ids`` use the +1 shifted
    slice so they align with the predicted positions.
    """
-    @property
+    required_keys = ["sequence", "loss_mask", "position_ids"]
    def required_keys(self) -> List[str]:
        return ["sequence", "loss_mask"]
-    def _fetch_data(self, begin_idx: int, end_idx: int, key: str) -> Tensor:
+    def __getitem__(self, index: int):
-        return self.storage.fetch(begin_idx, end_idx, key)
+        begin, end = self.store.sample_window(index)
-
+        x = self.store.fetch(begin, end, "sequence")
-    def __getitem__(self, index):
+        y = self.store.fetch(begin + 1, end + 1, "sequence")
-        begin_idx, end_idx = self.get_index(index)
+        position_ids = self.store.fetch(begin, end, "position_ids")
-
+        loss_mask = self.store.fetch(begin + 1, end + 1, "loss_mask")
-        x = self._fetch_data(begin_idx, end_idx, "sequence").to(dtype=torch.long)
+        return {
-        y = self._fetch_data(begin_idx + 1, end_idx + 1, "sequence").to(
+            "input_ids": x.to(dtype=torch.long),
-            dtype=torch.long
+            "target_ids": y.to(dtype=torch.long),
-        )
+            "position_ids": position_ids.to(dtype=torch.long),
-        loss_mask = self._fetch_data(begin_idx + 1, end_idx + 1, "loss_mask").to(
+            "loss_mask": loss_mask.to(dtype=torch.bool),
-            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."""
+    """Record-structured dataset for Direct Preference Optimization.
-    def __init__(self, window_size: int, stride: int):
+    Each sample is one preference pair (chosen + rejected) and is an
-        super().__init__(window_size, stride)
+    independent training unit — no windowing, stride, or cross-record
    concatenation.  This keeps each sequence self-contained so attention
    never leaks across preference pairs.
-    @property
+    Two loading paths (handled by :class:`DatasetFactory`):
    def required_keys(self) -> List[str]:
        return ["chosen", "rejected", "chosen_mask", "rejected_mask"]
-    def _fetch_data(self, begin_idx: int, end_idx: int, key: str) -> Tensor:
+    - **Pre-tokenized** (bin): ``store.load(path)`` reads per-record
-        return self.storage.fetch(begin_idx, end_idx, key)
+      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``.
    """
-    def __getitem__(self, index: int):
+    required_keys = ["chosen", "rejected", "chosen_mask", "rejected_mask"]
        begin_idx, end_idx = self.get_index(index)
-        chosen = self._fetch_data(begin_idx, end_idx, "chosen").to(dtype=torch.long)
+    def make_processor(self, tokenizer, max_len: int):
-        rejected = self._fetch_data(begin_idx, end_idx, "rejected").to(dtype=torch.long)
+        return partial(dpo_processor, tokenizer=tokenizer, max_len=max_len)
        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
        )
    def __getitem__(self, index: int) -> Dict[str, Tensor]:
        return {
-            "chosen": chosen,
+            "chosen": self.store.fetch_record(index, "chosen").to(dtype=torch.long),
-            "rejected": rejected,
+            "rejected": self.store.fetch_record(index, "rejected").to(dtype=torch.long),
-            "chosen_mask": chosen_mask,
+            "chosen_mask": self.store.fetch_record(index, "chosen_mask").to(
-            "rejected_mask": rejected_mask,
+                dtype=torch.bool
            ),
            "rejected_mask": self.store.fetch_record(index, "rejected_mask").to(
                dtype=torch.bool
            ),
        }
@DatasetFactory.register("grpo")
 class GRPODataset(BaseDataset):
-    """Dataset for Group Relative Policy Optimization training."""
+    """Dataset for offline Group Relative Policy Optimization.
-    def __init__(self, window_size: int, stride: int):
+    Each sample is one prompt with its group of responses and scalar
-        super().__init__(window_size, stride)
+    rewards — an independent training unit with no windowing or stride.
-    @property
+    Expected storage layout (produced by JsonlStore or pre-tokenized):
    def required_keys(self) -> List[str]:
        return ["prompts", "responses", "masks", "rewards"]
-    def _fetch_data(self, begin_idx: int, end_idx: int, key: str) -> Tensor:
+    - ``prompts``:   List[Tensor]  — one 1-D token tensor per record
-        return self.storage.fetch(begin_idx, end_idx, key)
+    - ``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]:
-        begin_idx, end_idx = self.get_index(index)
+        prompts = self.store.fetch_record(index, "prompts")
-
+        responses = self.store.fetch_record(index, "responses")
-        prompts = self._fetch_data(begin_idx, end_idx, "prompts").to(dtype=torch.long)
+        masks = self.store.fetch_record(index, "masks")
-        responses = self._fetch_data(begin_idx, end_idx, "responses").to(
+        rewards = self.store.fetch_record(index, "rewards")
            dtype=torch.long
        )
        masks = self._fetch_data(begin_idx, end_idx, "masks").to(dtype=torch.bool)
        rewards = self._fetch_data(begin_idx, end_idx, "rewards")
        return {
-            "prompts": prompts,
+            "prompts": prompts.to(dtype=torch.long),
-            "responses": responses,
+            "responses": [r.to(dtype=torch.long) for r in responses],
-            "masks": masks,
+            "masks": [m.to(dtype=torch.bool) for m in masks],
-            "rewards": rewards,
+            "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,
@@ -74,6 +82,7 @@ class ResumableDistributedSampler(Sampler[int]):
        self.epoch += 1
        self._indices = None
        self.iter = self.iter % self.num_samples_per_replica
    @property
    def _remaining(self):
@@ -1,97 +1,69 @@
 """Storage backends for different data formats.
-Layers:
+Architecture (composition over inheritance):
  - I/O layer:       save_* / load_* functions, read/write raw files (HDF5/bin)
                      return Dict[str, List[Tensor]] — format-specific, no state
  - Store (ABC):     central abstraction, normalizes multi-segment into
                      Dict[str, List[Tensor]] per key via _normalize(),
                      fetch() uses bisect across segments — no forced concat
  - Dataset layer:   BaseDataset owns a Store, only calls store.fetch(begin, end, key)
-Key properties:
+    Store (ABC)               — owns _data/_cum/_offsets bookkeeping
-  - Multi-segment:   segments kept as-is, no forced concatenation — safe for
+                                  + window_size/stride for sample-id
-                      datasets larger than RAM
+                                  indexing.  __getitem__/__len__ produce
-  - Explicit length: _length = min(total elements across keys), set at load,
+                                  the smallest iterable unit so Dataset
-                      __len__ returns O(1)
+                                  classes are pure delegators.
-  - Zero-copy mmap:  MmapStore wraps np.memmap(mode="r"), all DataLoader
+    Streamable (mixin)        — raw token slice  fetch(begin, end, keys)
-                      workers share OS page-cache pages
+    Recordable (mixin)        — raw record slice fetch_record(idx, keys)
    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 (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 os
+import logging
 from abc import ABC, abstractmethod
 from pathlib import Path
-from typing import Dict, List, Union
+from typing import Callable, Dict, List, Optional, Tuple, Union
 import h5py
 import numpy as np
 import torch
 from torch import Tensor
 from astrai.config.preprocess_config import PipelineConfig
 from astrai.factory import BaseFactory
 from astrai.preprocessing.transform import TokenizeTransform
 from astrai.serialization import (
    load_bin,
    load_bin_offsets,
 )
-
+logger = logging.getLogger(__name__)
 def save_h5(file_path: str, file_name: str, tensor_group: Dict[str, List[Tensor]]):
    os.makedirs(file_path, exist_ok=True)
    full_file_path = os.path.join(file_path, f"{file_name}.h5")
    with h5py.File(full_file_path, "w") as f:
        for key, tensors in tensor_group.items():
            grp = f.create_group(key)
            for idx, tensor in enumerate(tensors):
                arr = tensor.cpu().numpy()
                grp.create_dataset(f"data_{idx}", data=arr)
 def load_h5(file_path: str, share_memory=True) -> Dict[str, List[Tensor]]:
    tensor_group: Dict[str, List[Tensor]] = {}
    root_path = Path(file_path)
    h5_files = list(root_path.rglob("*.h5")) + list(root_path.rglob("*.hdf5"))
    for h5_file in h5_files:
        with h5py.File(h5_file, "r") as f:
            for key in f.keys():
                grp = f[key]
                dsets = []
                for dset_name in grp.keys():
                    dset = grp[dset_name]
                    tensor = torch.from_numpy(dset[:])
                    if share_memory:
                        tensor = tensor.share_memory_()
                    dsets.append(tensor)
                if tensor_group.get(key) is None:
                    tensor_group[key] = []
                tensor_group[key].extend(dsets)
    return tensor_group
 def save_bin(file_path: str, tensor_group: Dict[str, List[Tensor]]):
    os.makedirs(file_path, exist_ok=True)
    meta = {}
    for key, tensors in tensor_group.items():
        cat = torch.cat(tensors, dim=0)
        meta[key] = {"shape": list(cat.shape), "dtype": str(cat.dtype).split(".")[-1]}
        np.asarray(cat.cpu().numpy()).tofile(os.path.join(file_path, f"{key}.bin"))
    with open(os.path.join(file_path, "meta.json"), "w") as f:
        json.dump(meta, f)
 def load_bin(file_path: str) -> Dict[str, List[Tensor]]:
    with open(os.path.join(file_path, "meta.json"), "r") as f:
        meta = json.load(f)
    segments: Dict[str, List[Tensor]] = {}
    for key, info in meta.items():
        arr = np.memmap(
            os.path.join(file_path, f"{key}.bin"),
            dtype=info["dtype"],
            mode="r+",
            shape=tuple(info["shape"]),
        )
        segments[key] = [torch.from_numpy(arr)]
    return segments
 def detect_format(load_path: str) -> str:
@@ -101,7 +73,7 @@ def detect_format(load_path: str) -> str:
        load_path: Directory or file path
    Returns:
-        Format string ("h5" or "bin")
+        Format string ("h5", "bin", "jsonl", or "processed")
    Raises:
        FileNotFoundError: If no supported data files are found
@@ -109,142 +81,546 @@ def detect_format(load_path: str) -> str:
    root = Path(load_path)
    if root.is_file():
        suffix = root.suffix.lower()
-        if suffix in (".h5", ".hdf5"):
+        if suffix == ".jsonl":
-            return "h5"
+            return "jsonl"
        raise ValueError(f"Unsupported file format: {suffix}")
-    h5_files = list(root.rglob("*.h5")) + list(root.rglob("*.hdf5"))
+    bin_files = [Path(p) for p in glob.glob(str(root / "**" / "*.bin"), recursive=True)]
-    if h5_files:
+    if bin_files:
-        return "h5"
+        has_meta = (root / "meta.json").exists() or len(
-    bin_files = list(root.rglob("*.bin"))
+            [Path(p) for p in glob.glob(str(root / "**" / "meta.json"), recursive=True)]
-    if bin_files and (root / "meta.json").exists():
+        ) > 0
-        return "bin"
+        if has_meta:
            return "bin"
    jsonl_files = [
        Path(p) for p in glob.glob(str(root / "**" / "*.jsonl"), recursive=True)
    ]
    if jsonl_files:
        return "jsonl"
    raise FileNotFoundError(f"No supported data files found at {load_path}")
 class Store(ABC):
-    """String keys -> segmented tensors with ``fetch(begin, end, keys)``.
+    """Common base for all storage backends.
-    Each key maps to one or more tensor segments (no forced concatenation).
+    A Store owns both its data layout AND its sample-id → token/record
-    ``len(store)`` returns ``self._length`` (explicit, O(1)), the minimum
+    index translation.  Datasets are thin wrappers that bind a Store
-    total element count across all keys.
+    to a particular train-type's key mapping; they never know about
    window/stride math.
-    Subclasses fill ``self._data`` and ``self._cum`` during ``load()``
+    Two iteration modes:
-    via ``_normalize()``.
+
    - **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.
    """
-    def __init__(self):
+    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) -> None:
+    def load(self, path: str, **kwargs) -> None:
        raise NotImplementedError
    @property
    def keys(self) -> List[str]:
        return list(self._data.keys())
-    def __len__(self) -> int:
+    @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 (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 (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 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 (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]],
    ):
-        if not self._data:
+        return _stream_fetch(self, begin, end, keys)
            raise RuntimeError("Store not loaded")
        if not (0 <= begin < self._length and 0 <= end <= self._length):
            raise ValueError(
                f"Index out of bounds: begin={begin}, end={end}, length={self._length}"
            )
        if isinstance(keys, str):
            return self._fetch_key(keys, begin, end)
        return {k: self._fetch_key(k, begin, end) for k in keys}
    def _fetch_key(self, key: str, begin: int, end: int) -> Tensor:
        """Fetch slice [begin, end) across potentially multiple segments."""
        segments = self._data[key]
        cum = self._cum[key]
        seg_start = bisect.bisect_right(cum, begin)
        seg_end = bisect.bisect_left(cum, end)
-        results = []
+def _stream_fetch(self, begin: int, end: int, keys: Union[str, List[str]]):
-        for i in range(seg_start, seg_end + 1):
+    if not getattr(self, "_data", None):
-            prev = cum[i - 1] if i > 0 else 0
+        raise RuntimeError("Store not loaded")
-            s = max(begin - prev, 0)
+    if not (0 <= begin < self._length and 0 <= end <= self._length):
-            e = min(end - prev, segments[i].shape[0])
+        raise ValueError(
-            results.append(segments[i][s:e])
+            f"Index out of bounds: begin={begin}, end={end}, length={self._length}"
        return results[0] if len(results) == 1 else torch.cat(results, dim=0)
    def _normalize(self, raw: Dict[str, List[Tensor]]):
        """Register segments and pre-compute cumulative lengths.
        Does NOT concatenate — segments are kept as-is to avoid OOM on
        large datasets.  Sets ``self._length`` to the minimum total
        element count across all keys.
        """
        for key, tensors in raw.items():
            self._data[key] = tensors
            cum = []
            total = 0
            for t in tensors:
                total += t.shape[0]
                cum.append(total)
            self._cum[key] = cum
        self._length = (
            min((cum[-1] if cum else 0) for cum in self._cum.values())
            if self._cum
            else 0
        )
    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.
+    """Factory for creating Store instances by type name."""
    Example::
        @StoreFactory.register("custom")
        class CustomStore(Store):
            ...
    """
    @classmethod
    def _validate_component(cls, store_cls: type):
        if not issubclass(store_cls, Store):
            raise TypeError(f"{store_cls.__name__} must inherit from Store")
@StoreFactory.register("h5")
 class H5Store(Store):
    """HDF5-based storage backend (pre-tokenized data)."""
    def load(self, path: str):
        self._normalize(load_h5(path))
@StoreFactory.register("bin")
-class MmapStore(Store):
+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.
-    Format on disk::
+    Supports both access modes:
-        data_root/
+    - **Stream**: always available via :meth:`fetch`.
-          meta.json          # {key: {shape, dtype}, ...}
+    - **Record** (``fetch_record(i, key)``): only when ``meta.json``
-          <key>.bin          # raw numpy array, one per key
+      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``.
    """
-    def load(self, path: str):
+    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 = []
-        raw = load_bin(path)
+        root = Path(path)
-        self._normalize(raw)
+        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 eager/lazy 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.
    Three ways to supply an eager transform (first match wins):
    - **Explicit** (``transform=``): caller-built
      :class:`TokenizeTransform` applied eagerly.
    - **Config file**: ``dataset_config.json`` alongside the ``*.jsonl``
      files — loaded via :meth:`TokenizeTransform.from_config_file`.
    - **Default messages** (``tokenizer_path=`` given, no config file):
      a built-in chatml config that tokenises the ``messages`` field,
      masking every role except ``assistant`` (loss on assistant only).
      Lets SFT/SEQ train straight from a chat-style JSONL directory
      without a hand-written config.
    Two tokenisation modes, selected at :meth:`load` time:
    - **Eager** (default): applies the transform 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
    _DEFAULT_MESSAGES_CONFIG = {
        "version": 1,
        "input": {
            "sections": [{"field": "messages", "action": "$role", "template": True}]
        },
        "mask": {"system": "mask", "user": "mask", "assistant": "train"},
        "mask_default": "mask",
        "output": {"position_ids_mode": "doc_reset"},
    }
    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 not None and config_path.exists():
                transform = TokenizeTransform.from_config_file(str(config_path))
            else:
                tokenizer_path = kwargs.get("tokenizer_path")
                if not tokenizer_path:
                    raise FileNotFoundError(
                        f"JSONL dataset config not found. Expected "
                        f"{self.CONFIG_NAME} alongside *.jsonl files, pass an "
                        f"explicit transform, pass processor= for lazy "
                        f"on-the-fly tokenisation, or pass tokenizer_path= to "
                        f"use the built-in messages config."
                    )
                config = PipelineConfig.from_dict(self._DEFAULT_MESSAGES_CONFIG)
                transform = TokenizeTransform(config, tokenizer_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,55 @@
 """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)
    - ``AttentionBackend`` — ABC for attention computation strategies
    - ``TorchNativeBackend`` — default SDPA backend with KV cache I/O
    - ``CudaBackend`` — CUDA kernel backend with paged decode + prefill
 Layout convention: all q/k/v are ``[batch, seq_len, n_heads, head_dim]``
 (blhd). Scale is always ``1/sqrt(head_dim)``.
 Each wrapper calls its compiled CUDA kernel directly. Fallback to torch
 SDPA is handled by the attention backend, not the wrapper functions.
 """
 from astrai.extension.attention_backend import (
    ATTN_BACKEND,
    AttentionBackend,
    AttentionBackendFactory,
    CudaBackend,
    FlashAttnBackend,
    TorchNativeBackend,
    attention,
    attn_backend,
    get_backend,
 )
 from astrai.extension.attention_ops import (
    TensorLayout,
    attn_decode,
    attn_paged_decode,
    attn_prefill,
 )
 from astrai.extension.loader import KERNEL_NAMES, is_available
 from astrai.extension.rotary_backend import apply_rotary_emb
 __all__ = [
    "ATTN_BACKEND",
    "AttentionBackend",
    "AttentionBackendFactory",
    "CudaBackend",
    "TorchNativeBackend",
    "FlashAttnBackend",
    "TensorLayout",
    "attention",
    "attn_backend",
    "get_backend",
    "attn_decode",
    "attn_paged_decode",
    "attn_prefill",
    "is_available",
    "KERNEL_NAMES",
    "apply_rotary_emb",
 ]
@@ -0,0 +1,671 @@
 """Attention backend abstraction with context-manager switching.
 The backend encapsulates KV cache I/O and attention computation. The
 attention module (GQA/MLA) keeps projections, rotary, QK-norm, gating,
 and output projection; the backend handles everything from "write K/V
 to cache" through "SDPA output".
 Usage — mirroring ``torch.nn.attention.sdpa_kernel``:
    from astrai.extension import attn_backend, ATTN_BACKEND
    with attn_backend(ATTN_BACKEND.TORCH_NATIVE):
        engine.generate("hello")
    # or with an instance:
    with attn_backend(TorchNativeBackend()):
        ...
    # or the shorthand (instance is itself a context manager):
    with TorchNativeBackend():
        ...
 Thread-safe via ``contextvars`` — each scheduler thread gets its own
 active backend. ``get_backend()`` returns the active one, falling back
 to a process-wide default (cuda > flash > torch, overridable via
 ``ASTR_BACKEND``).
 Layout convention: all q/k/v are ``[batch, seq_len, n_heads, head_dim]``
 (blhd). The backend returns ``[batch, seq_len, n_heads * head_dim]``.
 """
 import contextvars
 import enum
 import functools
 import importlib
 import os
 from abc import ABC, abstractmethod
 from contextlib import contextmanager
 from typing import TYPE_CHECKING, Optional, Union
 import torch
 import torch.nn.functional as F
 from torch import Tensor
 from astrai.extension.attention_ops import (
    attn_paged_decode,
    attn_paged_prefill,
 )
 from astrai.extension.loader import is_available
 from astrai.factory import BaseFactory
 if TYPE_CHECKING:
    from astrai.inference.core.cache import KVCache
 _current_backend: contextvars.ContextVar["AttentionBackend"] = contextvars.ContextVar(
    "attn_backend"
 )
@functools.lru_cache(maxsize=1)
 def flash_attn_available() -> bool:
    if not torch.cuda.is_available():
        return False
    fa = _get_flash_attn()
    if fa is None:
        return False
    try:
        major = int(fa.__version__.split(".")[0])
        cc = torch.cuda.get_device_capability()
        cc_num = cc[0] * 10 + cc[1]
    except Exception:
        major, cc_num = 0, 0
    if (major >= 3 and cc_num < 90) or (major < 3 and 0 < cc_num < 70):
        return False
    try:
        if not hasattr(fa, "flash_attn_func"):
            return False
        x = torch.zeros(1, 1, 1, 64, device="cuda", dtype=torch.bfloat16)
        out = fa.flash_attn_func(x, x, x, causal=True)
        return bool(torch.isfinite(out).all().item())
    except Exception:
        return False
@functools.lru_cache(maxsize=1)
 def _get_flash_attn():
    try:
        return importlib.import_module("flash_attn")
    except Exception:
        return None
 class ATTN_BACKEND(enum.Enum):
    """Backend selector enum, mirroring ``torch.nn.attention.SDPBackend``."""
    TORCH_NATIVE = "torch_native"
    CUDA = "cuda"
    FLASH = "flash"
 _default_backend: Optional["AttentionBackend"] = None
 def _priority_backends() -> list["AttentionBackend"]:
    """Available backends in priority order: cuda -> flash -> torch."""
    backends: list[AttentionBackend] = []
    if is_available("attn_paged_decode") and is_available("attn_paged_prefill"):
        backends.append(CudaBackend())
    if flash_attn_available():
        backends.append(FlashAttnBackend())
    backends.append(TorchNativeBackend())
    return backends
 def _backend_supports(
    backend: "AttentionBackend",
    q: Tensor,
    kv_cache: Optional["KVCache"],
    attn_mask: Optional[Tensor],
    is_causal: bool,
 ) -> bool:
    """Whether ``backend`` can run this attention call.
    The CUDA kernels are bf16-only, support head_dim in 32/64/128/256, and
    need a KV cache (decode/prefill); everything else falls back to torch.
    """
    if isinstance(backend, CudaBackend):
        return (
            kv_cache is not None
            and q.dtype == torch.bfloat16
            and q.size(-1) in (32, 64, 128, 256)
        )
    if isinstance(backend, FlashAttnBackend):
        if not flash_attn_available():
            return False
        if q.size(1) == 1 and kv_cache is not None:
            return True
        return not (attn_mask is not None and not is_causal)
    return True
 def _resolve_default_backend() -> "AttentionBackend":
    """Pick the highest-priority available backend (cuda -> flash -> torch).
    Set ``ASTR_BACKEND`` to override: ``ASTR_BACKEND=cuda``, ``torch_native``,
    or ``flash``.  The value is the registered name (same as the
    ``ATTN_BACKEND`` enum value).
    Resolved lazily on first ``get_backend()`` and cached.  Per-call
    capability fallback happens in ``attention()``, so the default is
    safe for training and fp32 models.
    """
    forced = os.environ.get("ASTR_BACKEND", "").strip().lower()
    if forced:
        try:
            return AttentionBackendFactory.create(forced)
        except (ValueError, RuntimeError):
            pass
    return _priority_backends()[0]
 def get_backend() -> "AttentionBackend":
    """Return the active backend for the current thread/context.
    Falls back to the highest-priority available backend (cuda -> flash ->
    torch_native) when no backend has been activated via ``with``.  Set
    ``ASTR_BACKEND`` to override the default.
    """
    try:
        return _current_backend.get()
    except LookupError:
        global _default_backend
        if _default_backend is None:
            _default_backend = _resolve_default_backend()
        return _default_backend
@contextmanager
 def attn_backend(backend: Union[str, ATTN_BACKEND, "AttentionBackend", type]):
    """Context manager to select an attention backend.
    Mirrors ``torch.nn.attention.sdpa_kernel``. Accepts an
    registered name, ``ATTN_BACKEND`` enum value, backend class, or instance.
    Examples::
        with attn_backend(ATTN_BACKEND.TORCH_NATIVE):
            ...
        with attn_backend(TorchNativeBackend):
            ...
        with attn_backend(TorchNativeBackend()):
            ...
    """
    if isinstance(backend, ATTN_BACKEND):
        instance = AttentionBackendFactory.create(backend.value)
    elif isinstance(backend, str):
        instance = AttentionBackendFactory.create(backend)
    elif isinstance(backend, type) and issubclass(backend, AttentionBackend):
        instance = backend()
    elif isinstance(backend, AttentionBackend):
        instance = backend
    else:
        raise TypeError(
            f"expected a registered name, ATTN_BACKEND, AttentionBackend type, "
            f"or instance, "
            f"got {type(backend).__name__}"
        )
    token = _current_backend.set(instance)
    try:
        yield instance
    finally:
        _current_backend.reset(token)
 def repeat_kv(x: Tensor, n_rep: int) -> Tensor:
    """Expand KV heads to match Q heads 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 _write_and_gather_kv(
    kv_cache: "KVCache",
    k: Tensor,
    v: Tensor,
    layer_id: int,
    q: Tensor,
    attn_mask: Optional[Tensor],
 ) -> tuple[Tensor, Tensor]:
    kv_cache.k_buffer[layer_id, kv_cache.out_cache_loc] = k
    kv_cache.v_buffer[layer_id, kv_cache.out_cache_loc] = v
    max_len = kv_cache.max_len
    indices = kv_cache.req_to_token[kv_cache.req_pool_indices, :max_len]
    if q.size(1) == 1 and attn_mask is not None and attn_mask.dim() == 4:
        pos_mask = attn_mask[:, 0, 0]
    else:
        pos_mask = (
            torch.arange(max_len, device=q.device)[None, :] < kv_cache.seq_lens[:, None]
        )
    indices = torch.where(pos_mask, indices, torch.zeros_like(indices))
    return kv_cache.k_buffer[layer_id, indices], kv_cache.v_buffer[layer_id, indices]
 def attention(
    q: Tensor,
    k: Tensor,
    v: Tensor,
    kv_cache: Optional["KVCache"] = None,
    layer_id: int = 0,
    attn_mask: Optional[Tensor] = None,
    is_causal: bool = False,
 ) -> Tensor:
    """Functional attention entry point — mirrors ``F.scaled_dot_product_attention``.
    Delegates to the active backend (set via ``with attn_backend(...)``).
    Handles KV cache I/O, GQA head expansion, and causal masking so the
    caller only needs to provide projected q/k/v.
    Args:
        q: [batch, q_len, n_heads, head_dim] (blhd)
        k: [batch, q_len, n_kv_heads, head_dim] (blhd)
        v: [batch, q_len, n_kv_heads, head_dim] (blhd)
        kv_cache: cache dataclass, or None for training (no cache).
        layer_id: transformer layer index for buffer access.
        attn_mask: pre-built attention mask (SDPA-compatible).
        is_causal: whether to apply causal masking.
    Returns:
        [batch, q_len, n_heads * head_dim]
    """
    backend = get_backend()
    if not _backend_supports(backend, q, kv_cache, attn_mask, is_causal):
        # The active backend cannot run this call (e.g. CUDA on a training /
        # fp32 / unsupported-head_dim input) — fall back to the highest-
        # priority backend that can, ending at torch SDPA.
        for candidate in _priority_backends():
            if isinstance(candidate, type(backend)):
                continue
            if _backend_supports(candidate, q, kv_cache, attn_mask, is_causal):
                backend = candidate
                break
    return backend.forward(q, k, v, kv_cache, layer_id, attn_mask, is_causal)
 class AttentionBackend(ABC):
    """Abstract base for attention computation strategies.
    Subclasses implement ``fwd_decode`` (q_len == 1, with cache) and
    ``fwd_prefill`` (q_len > 1, with or without cache). The public
    ``forward`` method dispatches based on q_len.
    Three equivalent ways to activate a backend::
        with attn_backend(ATTN_BACKEND.TORCH_NATIVE):  # enum
            ...
        with attn_backend(TorchNativeBackend):          # class
            ...
        with TorchNativeBackend():                       # instance
            ...
    """
    def __enter__(self) -> "AttentionBackend":
        self._token = _current_backend.set(self)
        return self
    def __exit__(self, *exc) -> None:
        _current_backend.reset(self._token)
    def forward(
        self,
        q: Tensor,
        k: Tensor,
        v: Tensor,
        kv_cache: Optional["KVCache"],
        layer_id: int,
        attn_mask: Optional[Tensor] = None,
        is_causal: bool = False,
    ) -> Tensor:
        """Dispatch to decode or extend based on q_len.
        Args:
            q: [batch, q_len, n_heads, head_dim]
            k: [batch, q_len, n_kv_heads, head_dim]
            v: [batch, q_len, n_kv_heads, head_dim]
            kv_cache: cache dataclass, or None for training (no cache).
            layer_id: transformer layer index for buffer access.
            attn_mask: pre-built attention mask compatible with SDPA.
            is_causal: whether to apply causal masking.
        Returns:
            [batch, q_len, n_heads * head_dim]
        """
        if kv_cache is not None and q.size(1) == 1:
            return self.fwd_decode(q, k, v, kv_cache, layer_id, attn_mask, is_causal)
        return self.fwd_prefill(q, k, v, kv_cache, layer_id, attn_mask, is_causal)
    @abstractmethod
    def fwd_decode(
        self,
        q: Tensor,
        k: Tensor,
        v: Tensor,
        kv_cache: Optional["KVCache"],
        layer_id: int,
        attn_mask: Optional[Tensor] = None,
        is_causal: bool = False,
    ) -> Tensor:
        """Single-token decode with KV cache."""
    @abstractmethod
    def fwd_prefill(
        self,
        q: Tensor,
        k: Tensor,
        v: Tensor,
        kv_cache: Optional["KVCache"],
        layer_id: int,
        attn_mask: Optional[Tensor] = None,
        is_causal: bool = False,
    ) -> Tensor:
        """Multi-token prefill or training forward."""
    @staticmethod
    def supports_graph() -> bool:
        """Return True if this backend supports CUDA-graph capture.
        Override in subclasses that can run under ``torch.cuda.graph``.
        Called on the *active* backend instance (or its class) — a cheap
        boolean check with no side-effects.
        """
        return False
 class AttentionBackendFactory(BaseFactory[AttentionBackend]):
    """Factory for registered attention backends."""
@AttentionBackendFactory.register(ATTN_BACKEND.TORCH_NATIVE.value)
 class TorchNativeBackend(AttentionBackend):
    """Reference backend using torch SDPA with indirect KV cache indexing.
    Writes new K/V into the cache buffers, gathers the full sequence K/V
    via ``req_to_token`` indirect indexing, then calls
    ``F.scaled_dot_product_attention``.
    For training (``kv_cache is None``), skips cache I/O entirely and
    runs SDPA directly on the projected q/k/v.
    """
    @staticmethod
    def supports(**kwargs) -> bool:
        return True
    def fwd_decode(
        self,
        q: Tensor,
        k: Tensor,
        v: Tensor,
        kv_cache: Optional["KVCache"],
        layer_id: int,
        attn_mask: Optional[Tensor] = None,
        is_causal: bool = False,
    ) -> Tensor:
        return self._forward(q, k, v, kv_cache, layer_id, attn_mask, is_causal)
    def fwd_prefill(
        self,
        q: Tensor,
        k: Tensor,
        v: Tensor,
        kv_cache: Optional["KVCache"],
        layer_id: int,
        attn_mask: Optional[Tensor] = None,
        is_causal: bool = False,
    ) -> Tensor:
        return self._forward(q, k, v, kv_cache, layer_id, attn_mask, is_causal)
    def _forward(
        self,
        q: Tensor,
        k: Tensor,
        v: Tensor,
        kv_cache: Optional["KVCache"],
        layer_id: int,
        attn_mask: Optional[Tensor] = None,
        is_causal: bool = False,
    ) -> Tensor:
        if kv_cache is not None:
            k, v = _write_and_gather_kv(kv_cache, k, v, layer_id, q, attn_mask)
        n_rep = q.size(2) // k.size(2)
        if n_rep > 1:
            k = repeat_kv(k, n_rep)
            v = repeat_kv(v, n_rep)
        out = F.scaled_dot_product_attention(
            q.permute(0, 2, 1, 3),
            k.permute(0, 2, 1, 3),
            v.permute(0, 2, 1, 3),
            attn_mask,
            is_causal=is_causal,
        )
        out = out.permute(0, 2, 1, 3).contiguous().flatten(2)
        return out
@AttentionBackendFactory.register(ATTN_BACKEND.CUDA.value)
 class CudaBackend(AttentionBackend):
    """CUDA kernel backend with direct KV cache access.
    Decode path: writes K/V to the flat pool, then calls
    ``attn_paged_decode`` with req_to_token + kv_indptr.
    Prefill path: writes K/V to the flat pool, then calls
    ``attn_paged_prefill`` with ragged-batch support via qo_indptr +
    kv_indptr.
    ``kv_cache is None`` (training) raises — the per-call fallback to
    torch SDPA for training / fp32 / unsupported head_dim happens in the
    ``attention()`` entry point.
    Raises ``RuntimeError`` if the required kernel is not available.
    """
    @staticmethod
    def supports(**kwargs) -> bool:
        head_dim = kwargs.get("head_dim", -1)
        return (
            torch.cuda.is_available()
            and head_dim in (32, 64, 128, 256)
            and is_available("attn_paged_decode")
            and is_available("attn_paged_prefill")
        )
    @staticmethod
    def supports_graph() -> bool:
        return True
    def fwd_decode(
        self,
        q: Tensor,
        k: Tensor,
        v: Tensor,
        kv_cache: Optional["KVCache"],
        layer_id: int,
        attn_mask: Optional[Tensor] = None,
        is_causal: bool = False,
    ) -> Tensor:
        if kv_cache is None:
            raise RuntimeError("CudaBackend does not support training (kv_cache=None)")
        loc = kv_cache.out_cache_loc[:, 0]
        kv_cache.k_buffer[layer_id].index_copy_(0, loc, k[:, 0])
        kv_cache.v_buffer[layer_id].index_copy_(0, loc, v[:, 0])
        q_3d = q.squeeze(1)
        kv_indptr = kv_cache.kv_indptr
        out = attn_paged_decode(
            q_3d,
            kv_cache.k_buffer[layer_id],
            kv_cache.v_buffer[layer_id],
            kv_cache.req_to_token,
            kv_cache.req_pool_indices,
            kv_indptr,
            kv_cache.max_len,
            is_causal=True,
            o_part_buf=kv_cache.decode_o_part,
            ml_part_buf=kv_cache.decode_ml_part,
            out_buf=kv_cache.decode_out,
        )
        return out.unsqueeze(1).flatten(2)
    def fwd_prefill(
        self,
        q: Tensor,
        k: Tensor,
        v: Tensor,
        kv_cache: Optional["KVCache"],
        layer_id: int,
        attn_mask: Optional[Tensor] = None,
        is_causal: bool = False,
    ) -> Tensor:
        if kv_cache is None:
            raise RuntimeError("CudaBackend does not support training (kv_cache=None)")
        loc = kv_cache.out_cache_loc.reshape(-1)
        kv_cache.k_buffer[layer_id].index_copy_(
            0, loc, k.reshape(-1, k.size(2), k.size(3))
        )
        kv_cache.v_buffer[layer_id].index_copy_(
            0, loc, v.reshape(-1, v.size(2), v.size(3))
        )
        b = q.size(0)
        q_len = q.size(1)
        kv_indptr = kv_cache.kv_indptr
        qo_indptr = kv_cache.qo_indptr
        q_flat = q.reshape(b * q_len, q.size(2), q.size(3))
        out = attn_paged_prefill(
            q_flat,
            kv_cache.k_buffer[layer_id],
            kv_cache.v_buffer[layer_id],
            kv_cache.req_to_token,
            kv_cache.req_pool_indices,
            kv_indptr,
            qo_indptr,
            attn_mask,
            q_len,
            is_causal=is_causal,
        )
        return out.reshape(b, q_len, q.size(2), q.size(3)).flatten(2)
@AttentionBackendFactory.register(ATTN_BACKEND.FLASH.value)
 class FlashAttnBackend(AttentionBackend):
    """FlashAttention backend via the optional ``flash-attn`` package.
    Decode (q_len=1, contiguous cache): uses ``flash_attn_with_kvcache``,
    which reads K/V directly from the flat pool via cache_batch_idx +
    cache_seqlens — no materialized KV gather.
    Prefill / non-contiguous decode: falls back to KV gather +
    ``flash_attn_func``.
    """
    @staticmethod
    def supports(**kwargs) -> bool:
        return flash_attn_available()
    def fwd_decode(
        self,
        q: Tensor,
        k: Tensor,
        v: Tensor,
        kv_cache: Optional["KVCache"],
        layer_id: int,
        attn_mask: Optional[Tensor] = None,
        is_causal: bool = False,
    ) -> Tensor:
        return self._forward(q, k, v, kv_cache, layer_id, attn_mask, is_causal)
    def fwd_prefill(
        self,
        q: Tensor,
        k: Tensor,
        v: Tensor,
        kv_cache: Optional["KVCache"],
        layer_id: int,
        attn_mask: Optional[Tensor] = None,
        is_causal: bool = False,
    ) -> Tensor:
        return self._forward(q, k, v, kv_cache, layer_id, attn_mask, is_causal)
    def _forward(
        self,
        q: Tensor,
        k: Tensor,
        v: Tensor,
        kv_cache: Optional["KVCache"],
        layer_id: int,
        attn_mask: Optional[Tensor] = None,
        is_causal: bool = False,
    ) -> Tensor:
        if kv_cache is not None:
            if q.size(1) == 1 and kv_cache.k_buffer.size(
                1
            ) == kv_cache.req_to_token.size(0) * kv_cache.req_to_token.size(1):
                return self._decode_with_kvcache(q, k, v, kv_cache, layer_id)
            k, v = _write_and_gather_kv(kv_cache, k, v, layer_id, q, attn_mask)
        n_rep = q.size(2) // k.size(2)
        if n_rep > 1:
            k = repeat_kv(k, n_rep)
            v = repeat_kv(v, n_rep)
        if attn_mask is not None and not is_causal:
            raise ValueError(
                "FlashAttnBackend does not support a custom attention mask; "
                "use a causal mask or select TorchNativeBackend."
            )
        fa = _get_flash_attn()
        if fa is None:
            raise RuntimeError(
                "FlashAttnBackend requires the optional 'flash-attn' package. "
                "Install with `pip install flash-attn`."
            )
        out = fa.flash_attn_func(
            q.contiguous(), k.contiguous(), v.contiguous(), causal=is_causal
        )
        return out.contiguous().flatten(2)
    def _decode_with_kvcache(
        self,
        q: Tensor,
        k: Tensor,
        v: Tensor,
        kv_cache: "KVCache",
        layer_id: int,
    ) -> Tensor:
        max_batch = kv_cache.req_to_token.size(0)
        max_seq = kv_cache.req_to_token.size(1)
        n_kv = k.size(2)
        k_cache = kv_cache.k_buffer[layer_id].view(max_batch, max_seq, n_kv, k.size(3))
        v_cache = kv_cache.v_buffer[layer_id].view(max_batch, max_seq, n_kv, v.size(3))
        fa = _get_flash_attn()
        out = fa.flash_attn_with_kvcache(
            q=q,
            k_cache=k_cache,
            v_cache=v_cache,
            k=k,
            v=v,
            cache_seqlens=(kv_cache.seq_lens - 1).to(torch.int32),
            cache_batch_idx=kv_cache.req_pool_indices.to(torch.int32),
            causal=True,
        )
        return out.flatten(2)
@@ -0,0 +1,194 @@
 """Attention kernel wrapper functions — one entry point per compiled kernel.
 Each wrapper calls its CUDA kernel directly. If the kernel is not
 available, raises ``RuntimeError``. Fallback to torch SDPA is the
 responsibility of the attention backend, not this module.
 Layout convention: all q/k/v are ``[batch, seq_len, n_heads, head_dim]``
 (blhd). Scale is always ``1/sqrt(head_dim)``.
 Interface (all functions):
    is_causal: True = causal mask; False = non-causal
    mask:      2D [batch, kv_len] or 3D [batch, q_len, kv_len] (bool, True=keep)
 """
 import enum
 from typing import Optional
 import torch
 from astrai.extension.loader import _available, _modules
 class TensorLayout(enum.IntEnum):
    """Q/K/V tensor layout, mirrors the C++ ``TensorLayout`` enum in ``attn_common.h``.
    Kernels internally operate on BHLD; BLHD inputs are transposed at entry.
    """
    BHLD = 0  # [batch, n_heads, seq_len, head_dim]
    BLHD = 1  # [batch, seq_len, n_heads, head_dim]
 def _check_available(name: str):
    if not _available.get(name):
        raise RuntimeError(
            f"CUDA kernel '{name}' is not available. "
            f"Build with CSRC_KERNELS=true or use a torch-native backend."
        )
 def attn_decode(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    mask: Optional[torch.Tensor] = None,
    is_causal: bool = False,
 ) -> torch.Tensor:
    """GQA decode attention (q_len == 1).
    Args:
        q: [batch, 1, n_heads, head_dim] (blhd, bf16)
        k: [batch, kv_len, n_kv_heads, head_dim] (blhd, bf16)
        v: [batch, kv_len, n_kv_heads, head_dim] (blhd, bf16)
        mask: 2D [batch, kv_len] or 3D [batch, 1, kv_len] (bool, True=keep)
        is_causal: apply causal mask
    Returns:
        [batch, 1, n_heads, head_dim] (blhd, bf16)
    """
    _check_available("attn_decode")
    causal_offset = (k.size(1) - 1) if is_causal else -1
    return _modules["attn_decode"].attn_decode(
        q, k, v, mask=mask, causal_offset=causal_offset, layout=TensorLayout.BLHD
    )
 def attn_prefill(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    mask: Optional[torch.Tensor] = None,
    is_causal: bool = False,
 ) -> torch.Tensor:
    """GQA prefill attention (q_len > 1).
    Args:
        q: [batch, q_len, n_heads, head_dim] (blhd, bf16)
        k: [batch, kv_len, n_kv_heads, head_dim] (blhd, bf16)
        v: [batch, kv_len, n_kv_heads, head_dim] (blhd, bf16)
        mask: 2D [batch, kv_len] or 3D [batch, q_len, kv_len] (bool, True=keep)
        is_causal: apply causal mask
    Returns:
        [batch, q_len, n_heads, head_dim] (blhd, bf16)
    """
    _check_available("attn_prefill")
    causal_offset = (k.size(1) - q.size(1)) if is_causal else -1
    return _modules["attn_prefill"].attn_prefill(
        q, k, v, mask=mask, causal_offset=causal_offset, layout=TensorLayout.BLHD
    )
 def attn_paged_decode(
    q: torch.Tensor,
    k_cache: torch.Tensor,
    v_cache: torch.Tensor,
    req_to_token: torch.Tensor,
    req_pool_indices: torch.Tensor,
    kv_indptr: torch.Tensor,
    max_seq_len: int,
    mask: Optional[torch.Tensor] = None,
    is_causal: bool = False,
    o_part_buf: Optional[torch.Tensor] = None,
    ml_part_buf: Optional[torch.Tensor] = None,
    out_buf: Optional[torch.Tensor] = None,
 ) -> torch.Tensor:
    """SGLang-style paged decode (q_len == 1, flat KV pool).
    Reads K/V directly from a flat pool [size, kv_head, head_dim] via
    req_to_token indirect indexing.  Each request has its own seq_len
    (from kv_indptr), eliminating padding waste.
    Args:
        q: [batch, n_heads, head_dim] (bf16, 3D — no seq dim)
        k_cache: [pool_size, n_kv_heads, head_dim] (bf16, flat)
        v_cache: same as k_cache
        req_to_token: [num_reqs, max_context_len] (int64) — token -> slot
        req_pool_indices: [batch] (int64) — rows into req_to_token
        kv_indptr: [batch+1] (int32) — prefix sum of per-request seq_lens
        max_seq_len: max per-request seq_len (Python int, for split computation)
        mask: 2D [batch, max_seq_len] (bool, True=keep) or None
        is_causal: apply causal mask
        o_part_buf: pre-allocated split-KV o partial buffer (workflow bypass)
        ml_part_buf: pre-allocated split-KV m/l buffer (workflow bypass)
        out_buf: pre-allocated output buffer [batch, n_heads, head_dim] (graph-safe)
    Returns:
        [batch, n_heads, head_dim] (bf16, 3D)
    """
    _check_available("attn_paged_decode")
    causal_offset = 0 if is_causal else -1
    return _modules["attn_paged_decode"].attn_paged_decode(
        q,
        k_cache,
        v_cache,
        req_to_token,
        req_pool_indices,
        kv_indptr,
        max_seq_len,
        mask=mask,
        causal_offset=causal_offset,
        o_part_buf=o_part_buf,
        ml_part_buf=ml_part_buf,
        out_buf=out_buf,
    )
 def attn_paged_prefill(
    q: torch.Tensor,
    k_cache: torch.Tensor,
    v_cache: torch.Tensor,
    req_to_token: torch.Tensor,
    req_pool_indices: torch.Tensor,
    kv_indptr: torch.Tensor,
    qo_indptr: torch.Tensor,
    mask: Optional[torch.Tensor] = None,
    max_q_len: int = 0,
    is_causal: bool = False,
 ) -> torch.Tensor:
    """SGLang-style paged prefill (ragged batch, flat KV pool).
    Reads K/V directly from a flat pool [size, kv_head, head_dim] via
    req_to_token.  Supports ragged batches: each request has its own
    q_len and kv_len, addressed via qo_indptr and kv_indptr.
    Args:
        q: [total_q, n_heads, head_dim] (bf16, 3D — flattened across requests)
        k_cache: [pool_size, n_kv_heads, head_dim] (bf16, flat)
        v_cache: same as k_cache
        req_to_token: [num_reqs, max_context_len] (int64)
        req_pool_indices: [batch] (int64)
        kv_indptr: [batch+1] (int32) — prefix sum of per-request kv_lens
        qo_indptr: [batch+1] (int32) — prefix sum of per-request q_lens
        mask: 4D [batch, 1, q_len, kv_len] (bool, True=keep) or None
        max_q_len: max per-request q_len (Python int, for grid computation)
        is_causal: apply causal mask
    Returns:
        [total_q, n_heads, head_dim] (bf16, 3D)
    """
    _check_available("attn_paged_prefill")
    causal_offset = 0 if is_causal else -1
    return _modules["attn_paged_prefill"].attn_paged_prefill(
        q,
        k_cache,
        v_cache,
        req_to_token,
        req_pool_indices,
        kv_indptr,
        qo_indptr,
        mask,
        max_q_len,
        causal_offset=causal_offset,
    )
@@ -0,0 +1 @@
 """Compiled CUDA kernel modules (``*.so``) live here, kept separate from Python source."""
@@ -0,0 +1,42 @@
 """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",
    "attn_paged_prefill",
    "rotary_emb",
 ]
 _available: dict[str, bool] = {}
 _modules: dict[str, object] = {}
 for _name in KERNEL_NAMES:
    try:
        _mod = importlib.import_module(f".lib.{_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,54 @@
 """Rotary embedding with auto-dispatch to CUDA kernel.
 Single entry point ``apply_rotary_emb(x, freqs_cis)`` — uses the fused
 CUDA kernel when available, falls back to torch complex multiply otherwise.
 Layout: x is [batch, seq_len, n_heads, head_dim] (bf16).
 freqs_cis is [batch, seq_len, dim/2, 2] (f32) — [cos, sin] pairs.
 """
 import torch
 from torch import Tensor
 from astrai.extension.loader import is_available
 _cache = {"available": None}
 def _cuda_available() -> bool:
    if _cache["available"] is None:
        _cache["available"] = is_available("rotary_emb")
    return _cache["available"]
 def _torch_apply(x: Tensor, freqs_cis: Tensor) -> Tensor:
    cos, sin = freqs_cis[..., 0], freqs_cis[..., 1]
    dtype = x.dtype
    x_ = x.float().reshape(*x.shape[:-1], -1, 2)
    x_complex = torch.view_as_complex(x_)
    freqs_cis_complex = torch.complex(cos, sin).unsqueeze(2)
    x_rotated = x_complex * freqs_cis_complex
    x_out = torch.view_as_real(x_rotated).flatten(-2)
    return x_out.to(dtype)
 def apply_rotary_emb(x: Tensor, freqs_cis: Tensor) -> Tensor:
    """Apply rotary embedding to x.
    Args:
        x: [batch, seq_len, n_heads, head_dim] (bf16)
        freqs_cis: [batch, seq_len, dim/2, 2] (f32) — [cos, sin] pairs
    Returns:
        [batch, seq_len, n_heads, head_dim] (bf16)
    """
    if (
        _cuda_available()
        and not torch.is_grad_enabled()
        and x.is_cuda
        and x.dtype == torch.bfloat16
    ):
        from astrai.extension.rotary_ops import rotary_emb as _cuda_rotary
        return _cuda_rotary(x, freqs_cis)
    return _torch_apply(x, freqs_cis)
@@ -0,0 +1,39 @@
 """Rotary embedding CUDA kernel wrapper.
 Calls the compiled CUDA kernel directly. If the kernel is not available,
 raises ``RuntimeError``. Fallback to torch complex multiply is the
 responsibility of ``astrai.extension.rotary_backend.apply_rotary_emb``.
 Layout: x is [batch, seq_len, n_heads, head_dim] (bf16, contiguous).
 freqs_cis is [batch, seq_len, head_dim/2, 2] (f32, contiguous) — [cos, sin] pairs.
 """
 import torch
 from astrai.extension.loader import _available, _modules
 def _check_available():
    if not _available.get("rotary_emb"):
        raise RuntimeError(
            "CUDA kernel 'rotary_emb' is not available. "
            "Build with CSRC_KERNELS=true or use the torch fallback."
        )
 def rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
    """Fused rotary embedding kernel.
    Args:
        x: [batch, seq_len, n_heads, head_dim] (bf16, contiguous)
        freqs_cis: [batch, seq_len, head_dim/2, 2] (f32, contiguous) — [cos, sin] pairs
    Returns:
        [batch, seq_len, n_heads, head_dim] (bf16)
    """
    _check_available()
    if not x.is_contiguous():
        x = x.contiguous()
    if not freqs_cis.is_contiguous():
        freqs_cis = freqs_cis.contiguous()
    return _modules["rotary_emb"].rotary_emb(x, freqs_cis)
@@ -1,149 +1,124 @@
-"""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,
    get_args,
    get_origin,
 )
 T = TypeVar("T")
-class Registry:
+def _resolve_base_type(
-    """Flexible registry for component classes with category and priority support.
+    arg: Union[Type, str, ForwardRef], factory_cls: type
 ) -> Optional[Type]:
    """Resolve the generic type-arg T to a concrete class.
-    This registry stores component classes with optional metadata (category, priority).
+    - Concrete class (``BaseFactory[MyBase]``): returned directly.
-    It provides methods for registration, retrieval, and listing with filtering.
+    - Forward reference (``BaseFactory["MyBase"]``): ``Base["X"]``
      produces a ``ForwardRef("X")`` at class-creation time.  We
      extract the name and evaluate it in the factory module's
      global namespace — the same mechanism ``typing.get_type_hints``
      uses internally.
    """
    if isinstance(arg, type):
        return arg
-    def __init__(self):
+    if isinstance(arg, str):
-        self._entries = {}  # name -> (component_cls, category, priority)
+        name = arg
    elif isinstance(arg, ForwardRef):
        name = arg.__forward_arg__
    else:
        return None
-    def register(
+    mod = sys.modules.get(factory_cls.__module__)
-        self,
+    if mod is None:
-        name: str,
+        return None
-        component_cls: Type,
+    try:
-        category: Optional[str] = None,
+        return eval(name, vars(mod))  # noqa: S307
-        priority: int = 0,
+    except NameError:
-    ):
+        return 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:
        """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]:
+def _validate_component(component_cls: Type, base: Optional[Type]) -> None:
-        """Get component class with its metadata."""
+    """Validate that *component_cls* inherits from *base*.
        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:
+    No-op when *base* is ``None`` (e.g. forward-ref resolution failed).
-        """Check if a name is registered."""
+    """
-        return name in self._entries
+    if base is not None and not issubclass(component_cls, base):
-
+        raise TypeError(f"{component_cls.__name__} must inherit from {base.__name__}")
    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 registration.
-    This base class provides a decorator-based registration pattern
+    Create a factory by subclassing with the desired base type::
    for creating extensible component factories.
-    Example usage:
+        class MyFactory(BaseFactory[MyBase]):
        class MyFactory(BaseFactory[MyBaseClass]):
            pass
    Register components with the ``register`` decorator::
        @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_base_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)
+            _validate_component(component_cls, cls._component_base)
-            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.
        Filters kwargs to match the component's __init__ signature,
        so components don't need to declare **kwargs just to absorb
        parameters meant for other components.
        Args:
            name: Registered name of the component
            *args: Positional arguments passed to component constructor
            **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):
+        component_cls = cls._entries.get(name)
        if component_cls 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)
        sig = inspect.signature(component_cls.__init__)
        has_var_kwargs = any(
            p.kind == inspect.Parameter.VAR_KEYWORD for p in sig.parameters.values()
@@ -157,70 +132,22 @@ class BaseFactory(ABC, Generic[T]):
            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]):
        """Validate that the component class is valid for this factory.
        Override this method in subclasses to add custom validation.
        Args:
            component_cls: Component class to validate
        Raises:
            TypeError: If the component class is invalid
        """
        pass
    @classmethod
    def get_component_class(cls, name: str) -> Type[T]:
-        """Get the registered component class by name without instantiating it.
+        """Get the registered component class without instantiating it."""
-
+        entry = cls._entries.get(name)
-        Args:
+        if entry is None:
            name: Registered name of the component
        Returns:
            The component class itself
        Raises:
            ValueError: If the component name is not registered
        """
        if not cls._registry.contains(name):
            raise ValueError(
-                f"Unknown component: '{name}'. "
+                f"Unknown component: '{name}'. Supported types: {sorted(cls._entries)}"
                f"Supported types: {sorted(cls._registry.list_names())}"
            )
-        return cls._registry.get(name)
+        return entry
    @classmethod
-    def list_registered(cls) -> list:
+    def list_registered(cls) -> List[str]:
-        """List all registered component names.
+        """List all registered component names."""
-
+        return sorted(cls._entries)
        Returns:
            List of registered component names
        """
        return cls._registry.list_names()
    @classmethod
    def is_registered(cls, name: str) -> bool:
-        """Check if a component name is registered.
+        """Check if a component name is registered."""
-
+        return name in cls._entries
        Args:
            name: Component name to check
        Returns:
            True if registered, False otherwise
        """
        return cls._registry.contains(name)
    @classmethod
    def list_by_category(cls, category: str) -> List[str]:
        """List registered component names in a category."""
        return cls._registry.list_by_category(category)
    @classmethod
    def list_by_priority(cls, reverse: bool = False) -> List[str]:
        """List registered component names sorted by priority."""
        return cls._registry.list_by_priority(reverse)
 __all__ = ["Registry", "BaseFactory"]
@@ -5,19 +5,24 @@ Layers:
  - api/:         HTTP orchestration (ProtocolHandler, server)
  - protocols/:   Response builders (OpenAI, Anthropic)
  - transport/:   SSE transport utilities
-  - engine.py:    Facade (InferenceEngine), Value Object (GenerationRequest)
+  - engine.py:    Facade (InferenceEngine)
-  - sample.py:    Strategy pattern (TemperatureStrategy, TopKStrategy, TopPStrategy)
+  - sample.py:    Strategy pattern (TemperatureStrategy, TopKStrategy, TopPStrategy, FrequencyPenaltyStrategy)
 """
 from astrai.inference.api import (
    AnthropicMessage,
    BaseToolParser,
    ChatCompletionRequest,
    ChatMessage,
    FunctionDef,
    GenContext,
    MessagesRequest,
    ProtocolHandler,
    SimpleJsonToolParser,
    StopChecker,
-    app,
+    ToolDef,
    ToolParserFactory,
    get_app,
    run_server,
 )
 from astrai.inference.api.anthropic import AnthropicResponseBuilder
@@ -28,19 +33,19 @@ from astrai.inference.core import (
    Executor,
    InferenceScheduler,
    KVCache,
-    KvcacheView,
+    KVStorage,
    PagePool,
-    PrefixCache,
+    RadixCache,
-    Storage,
+    ReqToTokenPool,
    Task,
    TaskManager,
    TaskStatus,
    TaskTable,
    page_hash,
 )
-from astrai.inference.engine import GenerationRequest, InferenceEngine
+from astrai.inference.engine import InferenceEngine
 from astrai.inference.sample import (
    BaseSamplingStrategy,
    FrequencyPenaltyStrategy,
    SamplingPipeline,
    TemperatureStrategy,
    TopKStrategy,
@@ -50,7 +55,6 @@ from astrai.inference.sample import (
 __all__ = [
    "InferenceEngine",
    "GenerationRequest",
    "InferenceScheduler",
    "Executor",
    "STOP",
@@ -59,27 +63,32 @@ __all__ = [
    "TaskStatus",
    "Allocator",
    "KVCache",
-    "KvcacheView",
+    "KVStorage",
    "PagePool",
-    "PrefixCache",
+    "RadixCache",
-    "Storage",
+    "ReqToTokenPool",
    "TaskTable",
    "page_hash",
    "sample",
    "BaseSamplingStrategy",
    "TemperatureStrategy",
    "TopKStrategy",
    "TopPStrategy",
    "FrequencyPenaltyStrategy",
    "SamplingPipeline",
    "ProtocolHandler",
    "StopChecker",
    "GenContext",
    "BaseToolParser",
    "SimpleJsonToolParser",
    "ToolParserFactory",
    "OpenAIResponseBuilder",
    "AnthropicResponseBuilder",
    "ChatMessage",
    "ChatCompletionRequest",
    "FunctionDef",
    "ToolDef",
    "AnthropicMessage",
    "MessagesRequest",
-    "app",
+    "get_app",
    "run_server",
 ]
@@ -1,23 +1,39 @@
-"""Inference API: protocol handler, stop checker, and FastAPI server."""
+"""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,
-    app,
+    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",
-    "app",
+    "get_app",
    "run_server",
 ]
@@ -1,5 +1,6 @@
 """Anthropic message completion response builder."""
 import time
 import uuid
 from typing import Any, Dict, List, Tuple, Union
@@ -39,9 +40,8 @@ class AnthropicResponseBuilder(ResponseBuilder):
        prompt = engine.tokenizer.apply_chat_template(messages, tokenize=False)
        ctx = GenContext(
            resp_id=f"msg_{uuid.uuid4().hex[:24]}",
-            created=0,
+            created=int(time.time()),
            model=request.model,
            prompt_tokens=0,
        )
        stop_sequences = getattr(request, "stop_sequences", None) or []
        return prompt, ctx, stop_sequences
@@ -72,15 +72,17 @@ class AnthropicResponseBuilder(ResponseBuilder):
            ),
        ]
-    def format_chunk(self, token: str) -> str:
+    def format_chunk(self, token: str, **kwargs) -> List[str]:
-        return sse_event(
+        return [
-            {
+            sse_event(
-                "type": "content_block_delta",
+                {
-                "index": 0,
+                    "type": "content_block_delta",
-                "delta": {"type": "text_delta", "text": token},
+                    "index": 0,
-            },
+                    "delta": {"type": "text_delta", "text": token},
-            event="content_block_delta",
+                },
-        )
+                event="content_block_delta",
            )
        ]
    def format_stream_end(self, ctx: GenContext, stop: StopInfo) -> List[str]:
        events: List[str] = []
@@ -1,7 +1,9 @@
 """OpenAI chat completion response builder."""
 import logging
 import time
 import uuid
-from typing import Any, Dict, List, Tuple
+from typing import Any, Dict, List, Optional, Tuple, Union
 from pydantic import BaseModel
@@ -11,24 +13,78 @@ from astrai.inference.api.protocol import (
    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]
-        prompt = engine.tokenizer.apply_chat_template(messages, tokenize=False)
+        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=0,
+            created=int(time.time()),
            model=self._model,
            prompt_tokens=0,
        )
        stop = request.stop
        stop_sequences = (
@@ -55,7 +111,82 @@ class OpenAIResponseBuilder(ResponseBuilder):
            )
        ]
-    def format_chunk(self, token: str) -> str:
+    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,
@@ -63,12 +194,19 @@ class OpenAIResponseBuilder(ResponseBuilder):
                "created": 0,
                "model": self._model,
                "choices": [
-                    {"index": 0, "delta": {"content": token}, "finish_reason": None}
+                    {
                        "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(
                {
@@ -76,7 +214,9 @@ class OpenAIResponseBuilder(ResponseBuilder):
                    "object": "chat.completion.chunk",
                    "created": ctx.created,
                    "model": self._model,
-                    "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
+                    "choices": [
                        {"index": 0, "delta": {}, "finish_reason": finish_reason}
                    ],
                }
            ),
            sse_event(
@@ -91,6 +231,32 @@ class OpenAIResponseBuilder(ResponseBuilder):
    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",
@@ -35,7 +35,7 @@ class GenContext:
    resp_id: str
    created: int
    model: str
-    prompt_tokens: int
+    prompt_tokens: int = 0
    completion_tokens: int = 0
@@ -64,7 +64,7 @@ class StopChecker:
 class ResponseBuilder(ABC):
    """Interface for protocol-specific response formatting.
-    A new protocol requires one concrete builder implementing 6 methods.
+    A new protocol requires one concrete builder implementing 5 methods.
    """
    @abstractmethod
@@ -78,8 +78,15 @@ class ResponseBuilder(ABC):
        """SSE events that open the stream."""
    @abstractmethod
-    def format_chunk(self, token: str) -> str:
+    def format_chunk(self, token: str, **kwargs) -> List[str]:
-        """SSE event for a single generated token."""
+        """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]:
@@ -118,6 +125,7 @@ class ProtocolHandler:
            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:
@@ -137,15 +145,25 @@ class ProtocolHandler:
            body = ""
            yielded = ""
            matched = None
            token_ids: List[int] = []
            async for token in agen:
                ctx.completion_tokens += 1
                body += token
                new_ids = self.engine.tokenizer.encode(token)
                token_ids.extend(new_ids)
                matched = checker.check(body)
                if matched:
                    break
-                yield self.builder.format_chunk(token)
+                ctx.completion_tokens += 1
                for event in self.builder.format_chunk(
                    token,
                    body=body,
                    current_token_ids=token_ids,
                    delta_token_ids=new_ids,
                ):
                    yield event
                yielded += token
            stop = StopInfo(matched=matched, body=body, yielded=yielded)
@@ -168,7 +186,6 @@ class ProtocolHandler:
        matched = None
        async for token in agen:
            ctx.completion_tokens += 1
            chunks.append(token)
            body += token
@@ -176,6 +193,8 @@ class ProtocolHandler:
            if matched:
                break
            ctx.completion_tokens += 1
        content = "".join(chunks)
        stop = StopInfo(matched=matched, body=body)
        return self.builder.format_response(ctx, content, stop)
@@ -3,6 +3,9 @@ OpenAI / Anthropic-compatible chat completion server backed by continuous-batchi
 Protocol-specific formatting is delegated to ``astrai.inference.protocol``.
 This module owns the FastAPI app, request/response schemas, and dependency wiring.
 ``app`` is lazily constructed — importing this module does NOT create a FastAPI instance.
 Use :func:`get_app` to access the singleton.
 """
 import logging
@@ -12,7 +15,7 @@ from typing import Any, Dict, List, Optional, Union
 import torch
 import uvicorn
-from fastapi import FastAPI, HTTPException
+from fastapi import APIRouter, FastAPI, HTTPException
 from pydantic import BaseModel, Field
 from astrai.inference.api.anthropic import AnthropicResponseBuilder
@@ -24,12 +27,25 @@ from astrai.tokenize import AutoTokenizer
 logger = logging.getLogger(__name__)
-_project_root = Path(__file__).parent.parent.parent
+_app_instance: Optional[FastAPI] = None
 class ChatMessage(BaseModel):
    role: str
-    content: str
+    content: Optional[str] = None
    tool_calls: Optional[List[Dict[str, Any]]] = None
    tool_call_id: Optional[str] = None
 class FunctionDef(BaseModel):
    name: str
    description: Optional[str] = None
    parameters: Optional[Dict[str, Any]] = None
 class ToolDef(BaseModel):
    type: str = "function"
    function: FunctionDef
 class ChatCompletionRequest(BaseModel):
@@ -48,6 +64,8 @@ class ChatCompletionRequest(BaseModel):
    frequency_penalty: Optional[float] = Field(default=0.0, ge=-2.0, le=2.0)
    logit_bias: Optional[Dict[int, float]] = None
    user: Optional[str] = None
    tools: Optional[List[ToolDef]] = None
    tool_choice: Optional[Union[str, Dict[str, Any]]] = "auto"
 class AnthropicMessage(BaseModel):
@@ -84,17 +102,16 @@ async def lifespan(app: FastAPI):
        logger.info("Inference engine shutdown complete")
-app = FastAPI(title="AstrAI Inference Server", version="0.2.0", lifespan=lifespan)
+router = APIRouter()
 def _create_engine(
-    param_path: Optional[Path] = None,
+    param_path: Path,
    device: str = "cuda",
    dtype: torch.dtype = torch.bfloat16,
    max_batch_size: int = 16,
    max_seq_len: Optional[int] = None,
 ) -> InferenceEngine:
    if param_path is None:
        param_path = _project_root / "params"
    if not param_path.exists():
        raise FileNotFoundError(f"Parameter directory not found: {param_path}")
@@ -107,39 +124,56 @@ def _create_engine(
        model=model,
        tokenizer=tokenizer,
        max_batch_size=max_batch_size,
        max_seq_len=max_seq_len,
    )
    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 = app.state.engine
+    engine = get_app().state.engine
    if engine is None:
        raise HTTPException(status_code=503, detail="Engine not initialized")
    return engine
-@app.get("/health")
+@router.get("/health")
 async def health():
    app = get_app()
    return {
        "status": "ok",
        "model_loaded": app.state.engine is not None,
    }
-@app.get("/stats")
+@router.get("/stats")
 async def get_stats():
    return _get_engine().get_stats()
-@app.post("/v1/chat/completions")
+@router.post("/v1/chat/completions")
 async def chat_completion(request: ChatCompletionRequest):
    engine = _get_engine()
    handler = ProtocolHandler(request, engine, OpenAIResponseBuilder())
    return await handler.handle()
-@app.post("/v1/messages")
+@router.post("/v1/messages")
 async def create_message(request: MessagesRequest):
    engine = _get_engine()
    handler = ProtocolHandler(request, engine, AnthropicResponseBuilder())
@@ -147,19 +181,22 @@ async def create_message(request: MessagesRequest):
 def run_server(
    param_path: Path,
    host: str = "0.0.0.0",
    port: int = 8000,
    reload: bool = False,
    device: str = "cuda",
    dtype: torch.dtype = torch.bfloat16,
    param_path: Optional[Path] = None,
    max_batch_size: int = 16,
    max_seq_len: Optional[int] = None,
 ):
    app = get_app()
    app.state.server_config = {
        "device": device,
        "dtype": dtype,
        "param_path": param_path,
        "max_batch_size": max_batch_size,
        "max_seq_len": max_seq_len,
    }
    uvicorn.run(
        app,
@@ -0,0 +1,339 @@
 """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.
    Args:
        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.
        Args:
            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
@@ -3,11 +3,10 @@
 from astrai.inference.core.cache import (
    Allocator,
    KVCache,
-    KvcacheView,
+    KVStorage,
    PagePool,
-    PrefixCache,
+    RadixCache,
-    Storage,
+    ReqToTokenPool,
    TaskTable,
    page_hash,
 )
 from astrai.inference.core.executor import Executor
@@ -17,11 +16,10 @@ from astrai.inference.core.task import STOP, Task, TaskManager, TaskStatus
 __all__ = [
    "Allocator",
    "KVCache",
-    "KvcacheView",
+    "KVStorage",
    "PagePool",
-    "PrefixCache",
+    "RadixCache",
-    "Storage",
+    "ReqToTokenPool",
    "TaskTable",
    "page_hash",
    "Executor",
    "InferenceScheduler",
@@ -1,15 +1,34 @@
 """KV cache architecture: three-layer separation (SGLang-inspired).
 Layer 1 — KVStorage:      flat token-level K/V buffers [n_layers, size, H, D]
 Layer 2 — ReqToTokenPool:  index table [req_idx, pos] → physical token slot
 Layer 3 — Allocator:       slot/page allocation with ref-counting and LRU
 PagePool orchestrates all three plus RadixCache (prefix addressing).
 KVCache is a pure dataclass passed to the model for direct buffer access.
 Two modes:
  - contiguous (default): pre-allocated per-request blocks, no dynamic alloc
  - paged: shared pool with on-demand allocation, prefix caching support
 """
 import threading
 from collections import OrderedDict
-from typing import Callable, Dict, List, Optional, Tuple
+from dataclasses import dataclass
 from typing import Callable, Dict, List, Optional
 import torch
 from torch import Tensor
 from astrai.inference.core.workspace import InferenceWorkspace
-def page_hash(token_ids: List[int], page_idx: int, page_size: int) -> int:
+
 def page_hash(
    token_ids: List[int], page_idx: int, page_size: int, parent_hash: int = 0
 ) -> int:
    start = page_idx * page_size
    end = min(start + page_size, len(token_ids))
-    h = 0
+    h = parent_hash
    for i in range(start, end):
        h = (h * 31 + token_ids[i]) & 0xFFFFFFFFFFFFFFFF
    return h
@@ -62,307 +81,558 @@ class Allocator:
    def touch(self, idx: int):
        with self._lock:
-            self._lru.move_to_end(idx)
+            if idx in self._lru:
                self._lru.move_to_end(idx)
-class PrefixCache:
+class RadixNode:
-    """Hash-based prefix matching: maps page hashes to physical page indices."""
+    """A page-aligned edge in the CPU-side prefix radix."""
    __slots__ = ("parent", "children", "page_idx", "tokens", "lock_ref")
    def __init__(self, parent=None, tokens=(), page_idx=None):
        self.parent = parent
        self.children: Dict[tuple, "RadixNode"] = {}
        self.page_idx = page_idx
        self.tokens = tuple(tokens)
        self.lock_ref = 0
 class RadixCache:
    """Page-granular radix prefix index with exact token matching."""
    def __init__(self, page_size: int):
        self._page_size = page_size
        self._root = RadixNode()
        self._page_to_node: Dict[int, RadixNode] = {}
        # Retained as an introspection-compatible map; matching never relies on
        # this lossy value.
        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)
+            node = self._page_to_node.pop(idx, None)
-            if h is not None:
+            self._page_to_hash.pop(idx, None)
-                self._hash_to_page.pop(h, None)
+            if node is None:
                return
            node.page_idx = None
            parent = node.parent
            if parent is not None:
                parent.children.pop(node.tokens, None)
    def has_page(self, idx: int) -> bool:
        with self._lock:
-            return idx in self._page_to_hash
+            return idx in self._page_to_node
    def lookup(self, token_ids: List[int]) -> List[int]:
        with self._lock:
            full_pages = len(token_ids) // self._page_size
            hits: List[int] = []
            node = self._root
            for i in range(full_pages):
-                h = page_hash(token_ids, i, self._page_size)
+                start = i * self._page_size
-                p = self._hash_to_page.get(h)
+                page_tokens = tuple(token_ids[start : start + self._page_size])
-                if p is None:
+                child = node.children.get(page_tokens)
                if child is None or child.page_idx is None:
                    break
-                hits.append(p)
+                hits.append(child.page_idx)
                node = child
            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)
+            full_pages = len(token_ids) // self._page_size
-            old_h = self._page_to_hash.pop(page_idx, None)
+            if logical_page_idx >= full_pages:
-            if old_h is not None:
+                return
-                self._hash_to_page.pop(old_h, None)
+            old = self._page_to_node.pop(page_idx, None)
-            self._page_to_hash[page_idx] = h
+            self._page_to_hash.pop(page_idx, None)
-            self._hash_to_page[h] = page_idx
+            if old is not None and old.parent is not None:
                old.parent.children.pop(old.tokens, None)
            node = self._root
            for i in range(logical_page_idx + 1):
                start = i * self._page_size
                page_tokens = tuple(token_ids[start : start + self._page_size])
                child = node.children.get(page_tokens)
                if child is None:
                    child = RadixNode(node, page_tokens)
                    node.children[page_tokens] = child
                node = child
            if node.page_idx is not None and node.page_idx != page_idx:
                replaced = node.page_idx
                self._page_to_node.pop(replaced, None)
                self._page_to_hash.pop(replaced, None)
            node.page_idx = page_idx
            self._page_to_node[page_idx] = node
            self._page_to_hash[page_idx] = page_hash(
                token_ids, logical_page_idx, self._page_size
            )
    def release(self, pages: List[int]) -> None:
        with self._lock:
            for page_idx in pages:
                node = self._page_to_node.get(page_idx)
                if node is not None and node.lock_ref:
                    node.lock_ref -= 1
 class ReqToTokenPool:
    """Maps [req_idx, pos] -> physical token slot in KV storage.
    Each row is one request; each column is a sequence position. The value
    at [req_idx, pos] is the flat index into the KV storage buffers.
    """
    def __init__(self, size: int, max_context_len: int, device: torch.device):
        self.size = size
        self.max_context_len = max_context_len
        self.req_to_token = torch.zeros(
            (size, max_context_len), dtype=torch.long, device=device
        )
        self.free_slots = list(range(size))
        self._lock = threading.Lock()
    def alloc(self, num_reqs: int) -> Optional[List[int]]:
        with self._lock:
            if num_reqs > len(self.free_slots):
                return None
            slots = self.free_slots[:num_reqs]
            self.free_slots = self.free_slots[num_reqs:]
            return slots
    def free(self, req_indices: List[int]):
        with self._lock:
            self.free_slots.extend(req_indices)
    def write(self, indices, values):
        self.req_to_token[indices] = values
 class KVStorage:
    """Token-level KV cache storage.
    Buffers: [n_layers, size, n_kv_heads, head_dim]. Each token occupies
    one slot indexed by ReqToTokenPool.
    """
    def __init__(
        self,
        size: int,
        n_layers: int,
        n_kv_heads: int,
        head_dim: int,
        device: torch.device,
        dtype: torch.dtype,
    ):
        self.size = size
        self.k_buffer = torch.empty(
            (n_layers, size, n_kv_heads, head_dim), device=device, dtype=dtype
        )
        self.v_buffer = torch.empty(
            (n_layers, size, n_kv_heads, head_dim), device=device, dtype=dtype
        )
    def get_key_buffer(self, layer_id: int) -> Tensor:
        return self.k_buffer[layer_id]
    def get_value_buffer(self, layer_id: int) -> Tensor:
        return self.v_buffer[layer_id]
    def set_kv_buffer(self, layer_id: int, loc: Tensor, k: Tensor, v: Tensor) -> None:
        self.k_buffer[layer_id, loc] = k
        self.v_buffer[layer_id, loc] = v
@dataclass
 class KVCache:
    """Pure data struct passed to model for KV cache I/O.
    The attention layer does raw buffer indexing — no methods, no abstraction.
    Attributes:
        k_buffer: [n_layers, size, n_kv_heads, head_dim]
        v_buffer: [n_layers, size, n_kv_heads, head_dim]
        req_to_token: [num_reqs, max_ctx_len] — index table
        req_pool_indices: [batch_size] — row indices into req_to_token
        seq_lens: [batch_size] — per-request total sequence lengths
        out_cache_loc: [batch, new_seq_len] or [batch, 1] — write indices
        max_len: max(seq_lens) as Python int — avoids GPU sync in decode
        kv_indptr: [batch+1] int32 — prefix sum of seq_lens, precomputed once
            per step so the attention backend avoids rebuilding it per layer.
        qo_indptr: [batch+1] int32 — prefill qo prefix-sum (None in decode)
        decode_o_part: split-KV o partial workspace (mirrors FlashInfer)
        decode_ml_part: split-KV m/l partial workspace (mirrors FlashInfer)
        decode_out: pre-allocated decode output buffer (graph-safe)
    """
    k_buffer: Tensor
    v_buffer: Tensor
    req_to_token: Tensor
    req_pool_indices: Tensor
    seq_lens: Tensor
    out_cache_loc: Tensor
    max_len: int = 0
    kv_indptr: Optional[Tensor] = None
    qo_indptr: Optional[Tensor] = None
    decode_o_part: Optional[Tensor] = None
    decode_ml_part: Optional[Tensor] = None
    decode_out: Optional[Tensor] = None
 class PagePool:
-    """Orchestrates allocator (page management) and PrefixCache (content addressing)."""
+    """Top-level KV cache manager.
-    def __init__(self, allocator: Allocator, prefix: PrefixCache):
+    Combines KVStorage + ReqToTokenPool + Allocator + RadixCache.
        self._alloc = allocator
        self._prefix = prefix
        self._alloc.on_evict = prefix.evict
-    @property
+    Args:
-    def allocator(self) -> Allocator:
+        n_layers: Number of transformer layers.
-        return self._alloc
+        n_kv_heads: Number of KV attention heads.
        head_dim: Dimension per head.
        max_batch_size: Maximum concurrent requests.
        max_seq_len: Maximum sequence length per request.
        device, dtype: Tensor device and dtype.
        page_size: Page size for paged mode (1 = token-level).
        n_tokens: Total token slots for paged mode. None = contiguous mode
            (pre-allocates max_batch_size * max_seq_len).
    """
-    @property
+    def __init__(
-    def prefix(self) -> PrefixCache:
+        self,
-        return self._prefix
+        n_layers: int,
        n_kv_heads: int,
        head_dim: int,
        max_batch_size: int,
        max_seq_len: int,
        device: torch.device,
        dtype: torch.dtype,
        page_size: int = 1,
        n_tokens: Optional[int] = None,
    ):
        self.page_size = page_size
        self.max_batch_size = max_batch_size
        self.max_seq_len = max_seq_len
        self.device = device
        self.dtype = dtype
        self.n_layers = n_layers
        self.n_kv_heads = n_kv_heads
        self.head_dim = head_dim
-    def alloc(self) -> int:
+        self.contiguous = n_tokens is None
-        return self._alloc.alloc()
+        if self.contiguous:
            self.n_tokens = max_batch_size * max_seq_len
        else:
            self.n_tokens = n_tokens
-    def free(self, idx: int):
+        self._storage = KVStorage(
-        keep = self._prefix.has_page(idx)
+            self.n_tokens, n_layers, n_kv_heads, head_dim, device, dtype
-        self._alloc.free(idx, keep_cached=keep)
+        )
-        if not keep:
+        self._req_pool = ReqToTokenPool(max_batch_size, max_seq_len, device)
            self._prefix.evict(idx)
-    def inc_ref(self, idx: int):
+        if self.contiguous:
-        self._alloc.inc_ref(idx)
+            for i in range(max_batch_size):
                self._req_pool.req_to_token[i] = torch.arange(
                    i * max_seq_len, (i + 1) * max_seq_len, device=device
                )
            self._alloc: Optional[Allocator] = None
            self._prefix: Optional[RadixCache] = None
        else:
            n_pages = self.n_tokens // page_size
            self._alloc = Allocator(n_pages)
            self._prefix = RadixCache(page_size) if page_size > 1 else None
            if self._prefix is not None:
                self._alloc.on_evict = self._prefix.evict
-    def lookup(self, token_ids: List[int]) -> List[int]:
+        self._task_req: Dict[str, int] = {}
-        hits = self._prefix.lookup(token_ids)
+        self._task_len: Dict[int, int] = {}
-        for p in hits:
+        self._task_cached: Dict[str, int] = {}
-            self._alloc.touch(p)
+        self._task_slots: Dict[str, List[int]] = {}
-        return hits
+        self._task_pages: Dict[str, List[int]] = {}
    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):
+        # Steady-state decode validation state: the ordered task set and its
-        with self._lock:
+        # Python seq_lens mirror.  When the same set advances every sequence
-            self._pages[task_id] = page_table
+        # by exactly one token per step, bind_tasks updates the stable
-            self._cached[task_id] = cached
+        # buffers in-place (+=1 / +=inc) instead of re-cumsumming.  Any
        # task-set change is a miss and rebuilds.
        self._bind_sig: Optional[tuple] = None
        self._bind_seq_lens: Optional[List[int]] = None
-    def get(self, task_id: str) -> List[int]:
+    # ---- task lifecycle ----
        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 KvcacheView:
    """Bundles Storage + page_table + total_len for attention layers."""
    def __init__(self, storage: Storage, page_table: Tensor, total_len: int = 0):
        self._storage = storage
        self._page_table = page_table
        self._total_len = total_len
    def write(self, layer_id: int, k: Tensor, v: Tensor):
        start_pos = self._total_len - k.size(1)
        self._storage.write(layer_id, self._page_table, start_pos, k, v)
    def gather(self, layer_id: int) -> Tuple[Tensor, Tensor]:
        return self._storage.gather(layer_id, self._page_table, self._total_len)
 class KVCache:
    """Facade: page management + KV-cache I/O for continuous batching."""
    def __init__(
        self,
        n_layers: int,
        n_pages: int,
        page_size: int,
        n_kv_heads: int,
        head_dim: int,
        device: torch.device,
        dtype: torch.dtype,
    ):
        self.page_size = page_size
        self._pool = PagePool(Allocator(n_pages), PrefixCache(page_size))
        self._table = TaskTable(page_size)
        self._storage = Storage(
            n_layers, n_pages, page_size, n_kv_heads, head_dim, device, dtype
        )
    def task_alloc(self, task_id: str, prompt_ids: List[int]) -> bool:
-        hits = self._pool.lookup(prompt_ids)
+        req_slots = self._req_pool.alloc(1)
-        cached = len(hits) * self.page_size
+        if req_slots is None:
-        for p in hits:
+            return False
-            self._pool.inc_ref(p)
+        req_idx = req_slots[0]
        self._task_req[task_id] = req_idx
-        remaining = len(prompt_ids) - cached
+        if self.contiguous:
-        n_new = (
+            self._task_len[req_idx] = len(prompt_ids)
-            (remaining + self.page_size - 1) // self.page_size if remaining > 0 else 0
+            self._task_cached[task_id] = 0
-        )
+            return True
-        new_pages: List[int] = []
+
-        if n_new > 0:
+        n_tokens_needed = len(prompt_ids)
-            for _ in range(n_new):
+        cached = 0
-                p = self._pool.alloc()
+
-                if p < 0:
+        if self._prefix is not None:
-                    for hp in hits:
+            hits = self._prefix.lookup(prompt_ids)
-                        self._pool.free(hp)
+            cached = len(hits) * self.page_size
-                    for np in new_pages:
+            for p in hits:
-                        self._pool.free(np)
+                self._alloc.inc_ref(p)
            self._task_pages[task_id] = list(hits)
            self._task_slots[task_id] = []
        else:
            self._task_pages[task_id] = []
            self._task_slots[task_id] = []
        remaining = n_tokens_needed - cached
        if remaining > 0:
            if self.page_size == 1:
                slots = self._alloc_tokens(remaining)
                if slots is None:
                    for p in self._task_pages[task_id]:
                        self._alloc.free(p)
                    self._req_pool.free([req_idx])
                    del self._task_req[task_id]
                    return False
-                new_pages.append(p)
+                self._task_slots[task_id] = slots
            else:
                n_new_pages = (remaining + self.page_size - 1) // self.page_size
                new_pages = []
                for _ in range(n_new_pages):
                    p = self._alloc.alloc()
                    if p < 0:
                        for hp in self._task_pages[task_id]:
                            self._alloc.free(hp)
                        for np_ in new_pages:
                            self._alloc.free(np_)
                        self._req_pool.free([req_idx])
                        del self._task_req[task_id]
                        return False
                    new_pages.append(p)
                self._task_pages[task_id].extend(new_pages)
-        self._table.set(task_id, hits + new_pages, cached)
+        self._write_req_to_token(task_id, prompt_ids, cached)
        self._task_len[req_idx] = len(prompt_ids)
        self._task_cached[task_id] = cached
        return True
    def task_free(self, task_id: str):
-        page_table, _ = self._table.pop(task_id)
+        req_idx = self._task_req.pop(task_id, None)
-        for idx in page_table:
+        if req_idx is None:
-            self._pool.free(idx)
+            return
        self._task_len.pop(req_idx, None)
        self._task_cached.pop(task_id, None)
        if not self.contiguous:
            if self._prefix is not None:
                for p in self._task_pages.get(task_id, []):
                    keep = self._prefix.has_page(p)
                    self._alloc.free(p, keep_cached=keep)
                    if not keep:
                        self._prefix.evict(p)
            else:
                for p in self._task_pages.get(task_id, []):
                    self._alloc.free(p)
            self._task_pages.pop(task_id, None)
            self._task_slots.pop(task_id, None)
        self._req_pool.free([req_idx])
    def task_extend(self, task_id: str, pos: int) -> bool:
-        page_table = self._table.get(task_id)
+        req_idx = self._task_req.get(task_id)
-        needed = (pos + 1 + self.page_size - 1) // self.page_size
+        if req_idx is None or pos >= self.max_seq_len:
-        while len(page_table) < needed:
+            return False
-            p = self._pool.alloc()
+
        # Paged mode must also claim a physical slot for the new token;
        # contiguous mode's block is pre-allocated so this is a no-op.
        if not self.contiguous and not self._extend_slot(task_id, req_idx, pos):
            return False
        self._task_len[req_idx] = pos + 1
        return True
    def _extend_slot(self, task_id: str, req_idx: int, pos: int) -> bool:
        """Allocate the physical slot for one extended token (paged mode)."""
        if self.page_size == 1:
            slots = self._alloc_tokens(1)
            if slots is None:
                return False
            self._task_slots.setdefault(task_id, []).extend(slots)
            self._req_pool.req_to_token[req_idx, pos] = slots[0]
            return True
        page_idx = pos // self.page_size
        existing = self._task_pages.get(task_id, [])
        if page_idx >= len(existing):
            p = self._alloc.alloc()
            if p < 0:
                return False
-            page_table.append(p)
+            existing.append(p)
            self._task_pages[task_id] = existing
        page_offset = pos % self.page_size
        page = existing[page_idx]
        token_slot = page * self.page_size + page_offset
        self._req_pool.req_to_token[req_idx, pos] = token_slot
        return True
    def task_cached(self, task_id: str) -> int:
-        return self._table.get_cached(task_id)
+        return self._task_cached.get(task_id, 0)
    def task_record_hashes(
        self, task_id: str, prompt_ids: List[int], start_logical_page: int = 0
    ):
-        page_table = self._table.get(task_id)
+        if self._prefix is None or self.contiguous:
            return
        pages = self._task_pages.get(task_id, [])
        full_pages = len(prompt_ids) // self.page_size
-        for i in range(start_logical_page, full_pages):
+        for i in range(start_logical_page, min(full_pages, len(pages))):
-            self._pool.record(page_table[i], prompt_ids, i)
+            self._prefix.record(pages[i], prompt_ids, i)
-    def make_table_tensor(self, task_ids: List[str], device: torch.device) -> Tensor:
+    def task_cacheable_ids(
-        return self._table.table_tensor(task_ids, device)
+        self, task_id: str, prompt_ids: List[int], output_ids: List[int]
    ):
        """Return the sequence whose KV entries are already materialized.
-    def bind(self, page_table: Tensor, total_len: int = 0) -> KvcacheView:
+        The first sampled output is produced by prompt prefill, and the last
-        return KvcacheView(self._storage, page_table, total_len)
+        sampled output has not been decoded into KV yet.  Therefore the cache
        can safely retain the prompt plus every output except the last one.
        """
        return list(prompt_ids) + list(output_ids[:-1])
    # ---- bind for forward ----
    def bind_tasks(
        self,
        task_ids: List[str],
        workspace: InferenceWorkspace,
        device: Optional[torch.device] = None,
        start_pos: Optional[int] = None,
    ) -> KVCache:
        if device is None:
            device = workspace.device
        req_indices = [self._task_req[tid] for tid in task_ids]
        # Per-request lengths come from the pool's own tracking (task_alloc
        # sets len(prompt_ids); task_extend sets pos+1), so callers need not
        # pass them.
        seq_lens = [self._task_len[req_idx] for req_idx in req_indices]
        b = len(task_ids)
        sig = tuple(task_ids)
        # Write into the caller's workspace buffers (fixed addresses, sized
        # to max_batch/max_seq at init) — the sole owner of the per-step
        # KV bind tensors.
        rpi_buf = workspace.req_pool_indices
        sl_buf = workspace.seq_lens
        kvp_buf = workspace.kv_indptr
        inc_buf = workspace.inc
        ocl_buf = workspace.out_cache_loc
        incremental = (
            start_pos is None
            and self._bind_sig is not None
            and self._bind_sig == sig
            and self._bind_seq_lens is not None
            and len(self._bind_seq_lens) == b
            and all(s == p + 1 for s, p in zip(seq_lens, self._bind_seq_lens))
        )
        if incremental:
            # Steady-state decode: advance the stable buffers in-place.
            # Normal-mode buffers keep ``+=`` legal regardless of whether
            # this runs inside ``torch.inference_mode()``.
            sl_buf[:b] += 1
            kvp_buf[: b + 1] += inc_buf[: b + 1]
            req_pool_indices = rpi_buf[:b]
            seq_lens_t = sl_buf[:b]
            kv_indptr = kvp_buf[: b + 1]
        else:
            # Cold path: fill the stable buffers from fresh host tensors.
            rpi_buf[:b].copy_(
                torch.tensor(req_indices, dtype=torch.long, device=device)
            )
            sl_buf[:b].copy_(torch.tensor(seq_lens, dtype=torch.long, device=device))
            kvp_buf[: b + 1].zero_()
            kvp_buf[1 : b + 1] = sl_buf[:b].cumsum(0).to(torch.int32)
            req_pool_indices = rpi_buf[:b]
            seq_lens_t = sl_buf[:b]
            kv_indptr = kvp_buf[: b + 1]
        self._bind_sig = sig
        self._bind_seq_lens = list(seq_lens)
        if start_pos is not None:
            seq_len = seq_lens[0]
            out_cache_loc = self._req_pool.req_to_token[
                req_pool_indices, start_pos:seq_len
            ]
            # Ragged query segmentation for the prefill kernel, computed once
            # (was rebuilt per layer in CudaBackend.fwd_prefill).
            q_len = seq_len - start_pos
            workspace.qo_indptr[: b + 1].copy_(
                torch.arange(b + 1, dtype=torch.int32, device=device) * q_len
            )
            qo_indptr = workspace.qo_indptr[: b + 1]
            decode_o_part, decode_ml_part = None, None
            decode_out = None
        else:
            write_pos = seq_lens_t - 1
            loc = self._req_pool.req_to_token[req_pool_indices, write_pos].unsqueeze(-1)
            ocl_buf[:b].copy_(loc)
            out_cache_loc = ocl_buf[:b]
            qo_indptr = None
            decode_o_part = getattr(workspace, "decode_o_part", None)
            decode_ml_part = getattr(workspace, "decode_ml_part", None)
            decode_out = getattr(workspace, "decode_out", None)
        return KVCache(
            k_buffer=self._storage.k_buffer,
            v_buffer=self._storage.v_buffer,
            req_to_token=self._req_pool.req_to_token,
            req_pool_indices=req_pool_indices,
            seq_lens=seq_lens_t,
            out_cache_loc=out_cache_loc,
            max_len=max(seq_lens),
            kv_indptr=kv_indptr,
            qo_indptr=qo_indptr,
            decode_o_part=decode_o_part,
            decode_ml_part=decode_ml_part,
            decode_out=decode_out,
        )
    # ---- internals ----
    def _alloc_tokens(self, n: int) -> Optional[List[int]]:
        if self.page_size != 1:
            raise RuntimeError("_alloc_tokens is for page_size=1 only")
        slots = []
        for _ in range(n):
            p = self._alloc.alloc()
            if p < 0:
                for s in slots:
                    self._alloc.free(s)
                return None
            slots.append(p)
        return slots
    def _write_req_to_token(self, task_id: str, prompt_ids: List[int], cached: int):
        req_idx = self._task_req[task_id]
        total = len(prompt_ids)
        if self.contiguous:
            return
        if self.page_size == 1:
            slots = self._task_slots.get(task_id, [])
            all_slots = slots[: total - cached]
            if all_slots:
                self._req_pool.req_to_token[req_idx, cached:total] = torch.tensor(
                    all_slots, dtype=torch.long, device=self.device
                )
        else:
            pages = self._task_pages.get(task_id, [])
            for pos in range(cached, total):
                page_idx = pos // self.page_size
                page_offset = pos % self.page_size
                if page_idx < len(pages):
                    token_slot = pages[page_idx] * self.page_size + page_offset
                    self._req_pool.req_to_token[req_idx, pos] = token_slot
@@ -1,15 +1,151 @@
 import logging
 import os
 import time
 from contextlib import contextmanager
 from dataclasses import dataclass
 from typing import List, Optional
 import torch
 from torch import Tensor
-from astrai.inference.core.cache import KVCache
+from astrai.extension.attention_backend import (
    ATTN_BACKEND,
    CudaBackend,
    attn_backend,
    get_backend,
 )
 from astrai.inference.core.cache import PagePool
 from astrai.inference.core.graph import CudaGraphContext
 from astrai.inference.core.task import Task
 from astrai.inference.core.workspace import InferenceWorkspace
 from astrai.inference.sample import sample
 from astrai.model.automodel import AutoModel
 from astrai.tokenize.tokenizer import AutoTokenizer
 logger = logging.getLogger(__name__)
 _TIMED = os.environ.get("ASTRAI_TIMED", "") == "1"
@contextmanager
 def timed(label: str, log: Optional[logging.Logger] = None):
    """Wall-clock debug timer, enabled via ``ASTRAI_TIMED=1``."""
    if not _TIMED:
        yield
        return
    tic = time.perf_counter()
    yield
    elapsed_ms = (time.perf_counter() - tic) * 1000
    (log or logger).info("%s %.1fms", label, elapsed_ms)
@dataclass
 class SamplingBatchInfo:
    """Per-batch sampling parameters, cached across decode steps.
    Sampling params are constant for a given ordered task set, so they are
    built once (pinned-memory async H2D) and reused until the task set
    changes.  ``top_ks`` is int32 to match the native consumers.
    """
    temperatures: Tensor  # float32 [B]
    top_ks: Tensor  # int32  [B]
    top_ps: Tensor  # float32 [B]
    freq_penalties: Tensor  # float32 [B]
    has_freq: bool  # any frequency_penalty != 0 (avoids per-step GPU .any())
 def _build_sampling_batch_info(tasks: List[Task], device) -> SamplingBatchInfo:
    pin = str(device).startswith("cuda")
    freq_penalties = torch.tensor(
        [t.frequency_penalty for t in tasks], dtype=torch.float32, pin_memory=pin
    ).to(device, non_blocking=True)
    return SamplingBatchInfo(
        temperatures=torch.tensor(
            [t.temperature for t in tasks], dtype=torch.float32, pin_memory=pin
        ).to(device, non_blocking=True),
        top_ks=torch.tensor(
            [t.top_k for t in tasks], dtype=torch.int32, pin_memory=pin
        ).to(device, non_blocking=True),
        top_ps=torch.tensor(
            [t.top_p for t in tasks], dtype=torch.float32, pin_memory=pin
        ).to(device, non_blocking=True),
        freq_penalties=freq_penalties,
        has_freq=bool((freq_penalties != 0).any()),
    )
 def _warmup_cuda_graphs(
    model: AutoModel,
    pool: PagePool,
    ws: InferenceWorkspace,
    gctx: "CudaGraphContext",
    max_batch_size: int,
    prompt_len: int = 1,
    device: Optional[str] = None,
 ):
    batch_sizes = [1]
    n = 2
    while n <= max_batch_size:
        batch_sizes.append(n)
        n *= 2
    if max_batch_size not in batch_sizes:
        batch_sizes.append(max_batch_size)
    dev = device or next(model.parameters()).device
    for b in batch_sizes:
        task_ids = [f"_gr_{b}_{i}" for i in range(b)]
        prompt_tokens = [list(range(prompt_len)) for _ in range(b)]
        alloc_ok = True
        for tid, pt in zip(task_ids, prompt_tokens):
            if not pool.task_alloc(tid, pt):
                alloc_ok = False
                break
        if not alloc_ok:
            for tid in task_ids:
                pool.task_free(tid)
            continue
        with (
            torch.inference_mode(),
            attn_backend(ATTN_BACKEND.CUDA),
            timed(f"warmup prefill b={b}", logger),
        ):
            kv_cache = pool.bind_tasks(task_ids, ws, start_pos=0)
            ids_in = torch.tensor(prompt_tokens, dtype=torch.long, device=dev)
            pos_in = torch.arange(prompt_len, device=dev).unsqueeze(0).expand(b, -1)
            model(
                ids_in,
                input_mask=pos_in.unsqueeze(-1) >= torch.arange(prompt_len, device=dev),
                kv_cache=kv_cache,
                position_ids=pos_in,
            )
        with (
            torch.inference_mode(),
            attn_backend(ATTN_BACKEND.CUDA),
            timed(f"warmup decode b={b}", logger),
        ):
            for step in range(2):
                seq_pos = prompt_len + step
                ws.position_ids[:b] = seq_pos
                for tid in task_ids:
                    pool.task_extend(tid, seq_pos)
                kv = pool.bind_tasks(task_ids, ws)
                input_mask = ws.decode_mask(ws.position_ids[:b], ws.max_seq_len)
                ids_buf = ws.fill_input_ids([step] * b)
                gctx.forward(
                    model,
                    key=(b,),
                    input_ids=ids_buf.unsqueeze(1),
                    input_mask=input_mask,
                    kv_cache=kv,
                    position_ids=ws.position_ids[:b].unsqueeze(1),
                )
        for tid in task_ids:
            pool.task_free(tid)
        torch.cuda.synchronize()
 class Executor:
@@ -19,20 +155,120 @@ class Executor:
        self,
        model: AutoModel,
        tokenizer: AutoTokenizer,
-        page_cache: KVCache,
+        kv_cache: PagePool,
        device: Optional[str] = None,
        dtype: Optional[torch.dtype] = None,
    ):
        self.model = model
        self.tokenizer = tokenizer
-        self.page_cache = page_cache
+        self.kv_cache = kv_cache
        self.device = device or next(model.parameters()).device
        self.dtype = dtype or next(model.parameters()).dtype
-    def execute_prefill(self, tasks: List[Task], prompt_len: int, start_pos: int = 0):
+        # Per-step decode cache for the steady-state case where the same
-        if start_pos >= prompt_len:
+        # ordered task set decodes one token per step.  Sampling params are
        # constant across steps; position_ids grows by exactly 1.  Single-slot:
        # any task-set change is a cache miss.
        self._decode_cache: Optional[tuple] = None
        # Pre-allocated fixed-shape buffers for the decode hot path
        # (input_ids, decode mask, KV bind metadata).  Eagerly sized at init
        # so the workspace is CUDA-graph-capture friendly — no allocation
        # during capture.
        config = model.config
        max_q_heads = config.num_attention_heads
        head_dim = config.hidden_size // config.num_attention_heads
        self._head_dim = head_dim
        self._graph_supported = CudaBackend.supports(head_dim=head_dim)
        self._workspace = InferenceWorkspace(
            max_batch_size=kv_cache.max_batch_size,
            max_seq_len=kv_cache.max_seq_len,
            max_q_heads=max_q_heads,
            head_dim=head_dim,
            device=self.device,
            dtype=self.dtype,
        )
        # CUDA-graph capture: one graph per (batch_size,) key.
        # Enabled at init-time via _warmup_cuda_graphs for CudaBackend
        # on supported head_dims; left disabled otherwise.
        self._graph_ctx = CudaGraphContext()
        self._try_enable_cuda_graph()
    def _try_enable_cuda_graph(self):
        if not self._graph_supported:
            return
        self._graph_ctx.set_enabled(True)
        _warmup_cuda_graphs(
            self.model,
            self.kv_cache,
            self._workspace,
            self._graph_ctx,
            max_batch_size=self.kv_cache.max_batch_size,
            device=self.device,
        )
    def _sample_logits(
        self,
        logits: Tensor,
        tasks: List[Task],
        return_logprobs: bool = False,
        info: Optional[SamplingBatchInfo] = None,
    ):
        info = info or _build_sampling_batch_info(tasks, self.device)
        if info.has_freq:
            history_lists = [
                t.prompt_ids[-t.rep_window :] + t.output_ids for t in tasks
            ]
            history_lens = [len(ids) for ids in history_lists]
            max_len = max(history_lens, default=0)
            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, ids in enumerate(history_lists):
                length = len(ids)
                padded_ids[i, :length] = torch.as_tensor(
                    ids, dtype=torch.long, device=self.device
                )
                padded_mask[i, :length] = True
        else:
            padded_ids = None
            padded_mask = None
        result = sample(
            logits,
            temperature=info.temperatures,
            top_k=info.top_ks,
            top_p=info.top_ps,
            frequency_penalty=info.freq_penalties,
            input_ids=padded_ids,
            input_mask=padded_mask,
            return_logprobs=return_logprobs,
        )
        if not return_logprobs:
            return result.tolist()
        tokens, logprobs = result
        tokens_list = tokens.tolist()
        logprobs_list = logprobs.tolist()
        for task, logprob in zip(tasks, logprobs_list):
            task.output_logprobs.append(float(logprob))
        return list(zip(tokens_list, logprobs_list))
    def execute_prefill(
        self,
        tasks: List[Task],
        prompt_len: int,
        start_pos: int = 0,
        return_logprobs: bool = False,
    ):
        if start_pos >= prompt_len:
            return []
        tasks = sorted(tasks, key=lambda t: t.task_id)
        batch_sz = len(tasks)
@@ -43,52 +279,118 @@ class Executor:
        )
        task_ids = [t.task_id for t in tasks]
-        page_tables = self.page_cache.make_table_tensor(task_ids, self.device)
+        position_ids = (
-
+            torch.arange(start_pos, prompt_len, dtype=torch.long, device=self.device)
-        with torch.inference_mode():
+            .unsqueeze(0)
-            self.model(
+            .expand(batch_sz, -1)
-                input_ids,
+        )
-                position_ids=torch.arange(
+        input_mask = position_ids.unsqueeze(-1) >= torch.arange(
-                    start_pos, prompt_len, dtype=torch.long, device=self.device
+            prompt_len, device=self.device
                )
                .unsqueeze(0)
                .expand(batch_sz, -1),
                paged_cache=self.page_cache.bind(page_tables, total_len=prompt_len),
            )
    def execute_decode(self, tasks: List[Task]) -> List[int]:
        if not tasks:
            return []
        input_ids = torch.tensor(
            [t.output_ids[-1] if t.output_ids else t.prompt_ids[-1] for t in tasks],
            dtype=torch.long,
            device=self.device,
        )
-        position_ids = torch.tensor(
+        with (
-            [t.next_pos for t in tasks], dtype=torch.long, device=self.device
+            torch.inference_mode(),
-        )
+            timed(f"execute_prefill b={batch_sz} prompt_len={prompt_len}", logger),
-        total_len = position_ids.max().item() + 1
+        ):
        task_ids = [t.task_id for t in tasks]
        page_tables = self.page_cache.make_table_tensor(task_ids, self.device)
        temperatures = torch.tensor([t.temperature for t in tasks], device=self.device)
        top_ks = torch.tensor([t.top_k for t in tasks], device=self.device)
        top_ps = torch.tensor([t.top_p for t in tasks], device=self.device)
        with torch.inference_mode():
            outputs = self.model(
-                input_ids.unsqueeze(1),
+                input_ids,
-                paged_cache=self.page_cache.bind(page_tables, total_len=total_len),
+                input_mask=input_mask,
-                position_ids=position_ids.unsqueeze(1),
+                position_ids=position_ids,
                kv_cache=self.kv_cache.bind_tasks(
                    task_ids,
                    self._workspace,
                    start_pos=start_pos,
                ),
            )
            logits = outputs["logits"][:, -1, :]
-        return sample(
+        return tasks, self._sample_logits(logits, tasks, return_logprobs)
-            logits,
+
-            temperature=temperatures,
+    def execute_decode(
-            top_k=top_ks,
+        self, tasks: List[Task], return_logprobs: bool = False
-            top_p=top_ps,
+    ) -> List[int]:
-        ).tolist()
+        """Decode next token for each task.
        Args:
            return_logprobs: When ``True``, also record (and return)
                the log-probability of each sampled token under the
                post-strategy sampling distribution.  The logprob is
                appended to ``task.output_logprobs`` and the return
                list becomes ``List[Tuple[int, float]]``.
        Returns:
            ``List[int]`` of sampled token IDs, or
            ``List[Tuple[int, float]]`` of ``(token_id, logprob)`` when
            ``return_logprobs`` is ``True``.
        """
        if not tasks:
            return []
        b = len(tasks)
        ws = self._workspace
        # ---- pre-replay: update input buffers in-place ----
        input_ids = ws.fill_input_ids(
            [t.output_ids[-1] if t.output_ids else t.prompt_ids[-1] for t in tasks]
        )
        task_ids = [t.task_id for t in tasks]
        cur_positions = [t.next_pos for t in tasks]
        sig = tuple(task_ids)
        cached = self._decode_cache
        if (
            cached is not None
            and cached[0] == sig
            and cur_positions == [p + 1 for p in cached[1]]
        ):
            info = cached[2]
            ws.position_ids[:b] += 1
            self._decode_cache = (sig, cur_positions, info)
        else:
            info = _build_sampling_batch_info(tasks, self.device)
            ws.position_ids[:b].copy_(
                torch.tensor(cur_positions, dtype=torch.long, device=self.device)
            )
            self._decode_cache = (sig, cur_positions, info)
        total_len = max(cur_positions) + 1
        input_mask = ws.decode_mask(ws.position_ids[:b], total_len)
        kv_cache = self.kv_cache.bind_tasks(task_ids, ws)
        # ---- forward (graph replay or live run + capture) ----
        use_graph = (
            self._graph_ctx.enabled
            and self._graph_supported
            and get_backend().supports_graph()
        )
        key = (b,)
        if use_graph:
            input_mask = ws.decode_mask(ws.position_ids[:b], ws.max_seq_len)
        with (
            torch.inference_mode(),
            timed(f"execute_decode forward b={b}", logger),
        ):
            if use_graph:
                outputs = self._graph_ctx.forward(
                    self.model,
                    key=key,
                    input_ids=input_ids.unsqueeze(1),
                    input_mask=input_mask,
                    kv_cache=kv_cache,
                    position_ids=ws.position_ids[:b].unsqueeze(1),
                )
            else:
                outputs = self.model(
                    input_ids.unsqueeze(1),
                    input_mask=input_mask,
                    kv_cache=kv_cache,
                    position_ids=ws.position_ids[:b].unsqueeze(1),
                )
            logits = outputs["logits"][:, -1, :]
        return self._sample_logits(logits, tasks, return_logprobs, info=info)
@@ -0,0 +1,103 @@
 """CUDA-graph capture for the decode model-forward step.
 Mirrors SGLang's cuda-graph manager: one graph per batch size.  The graph
 pair.  The graph captures ``model.forward()`` with workspace-backed inputs
 (all at fixed addresses).  Before each replay the caller updates the input
 buffer content in-place so the graph sees fresh data at the same tensor
 addresses.
 Only the model forward is captured — sampling runs outside the graph
 (via ``torch.multinomial`` which consumes a mutable RNG state).
 """
 import torch
 from torch import Tensor
 class CudaGraphContext:
    """CUDA-graph capture/replay for decode steps.
    Parameters:
        enabled: When ``False``, ``forward()`` always runs the live model
            forward without capture/replay (graphs are cleared).  Toggle at
            runtime via the ``set_enabled()`` method.
    Usage::
        gctx = CudaGraphContext()
        with torch.inference_mode():
            outputs = gctx.forward(
                model,
                key=(batch_size,),
                input_ids=workspace.input_ids[:b].unsqueeze(1),
                input_mask=input_mask,
                kv_cache=kv_cache,
                position_ids=workspace.position_ids[:b].unsqueeze(1),
            )
    The first call at a given key runs *without* capture (warmup).  The
    second call captures the graph.  Subsequent calls replay the captured
    graph.  A ``torch.cuda.synchronize()`` before capture drains in-flight
    work so the graph trace is clean.
    """
    def __init__(self, enabled: bool = False):
        self._enabled = enabled
        self._graphs: dict[tuple, torch.cuda.CUDAGraph] = {}
        self._outputs: dict[tuple, dict[str, Tensor]] = {}
        self._warmed: set[tuple] = set()
    @property
    def enabled(self) -> bool:
        return self._enabled
    def set_enabled(self, flag: bool):
        """Enable or disable CUDA-graph capture at runtime.
        Disabling clears all captured graphs (frees GPU memory) and warmup
        state.  Re-enabling after disable starts fresh — graphs are
        re-captured on the next warmup cycle.
        """
        if flag == self._enabled:
            return
        self._enabled = flag
        if not flag:
            self._graphs.clear()
            self._outputs.clear()
            self._warmed.clear()
    def forward(self, model, *, key, **kwargs) -> dict[str, Tensor]:
        """Run ``model(**kwargs)`` via graph replay or live forward.
        Args:
            model: callable, e.g. ``self.model.forward``.
            key: ``(batch_size,)`` — the dispatch key (one graph per batch size).
            **kwargs: arguments forwarded to ``model``.  All tensor arguments
                must reside at stable addresses (workspace buffers).
        Returns:
            The dict produced by ``model(**kwargs)``, e.g.
            ``{"logits": ..., "h0": ...}``.
        """
        if not self._enabled:
            self._outputs[key] = model(**kwargs)
            return self._outputs[key]
        if key in self._graphs:
            self._graphs[key].replay()
        elif key in self._warmed:
            cap_output = model(**kwargs)
            torch.cuda.synchronize()
            graph = torch.cuda.CUDAGraph()
            with torch.cuda.graph(graph):
                self._outputs[key] = model(**kwargs)
            self._graphs[key] = graph
            self._warmed.discard(key)
            return cap_output
        else:
            self._warmed.add(key)
            self._outputs[key] = model(**kwargs)
        return self._outputs[key]
    def has_graph(self, key: tuple) -> bool:
        return key in self._graphs
@@ -1,10 +1,11 @@
 import logging
 import threading
 import uuid
 from typing import Any, Dict, List, Optional, Tuple
 import torch
-from astrai.inference.core.cache import KVCache
+from astrai.inference.core.cache import PagePool
 from astrai.inference.core.executor import Executor
 from astrai.inference.core.task import STOP, Task, TaskManager, TaskStatus
 from astrai.model.automodel import AutoModel
@@ -14,7 +15,7 @@ logger = logging.getLogger(__name__)
 class InferenceScheduler:
-    """Four-phase continuous batching loop: cleanup -> refill -> prefill -> decode."""
+    """Continuous batching loop: cleanup -> refill -> prefill -> decode (all groups)."""
    def __init__(
        self,
@@ -22,73 +23,155 @@ class InferenceScheduler:
        tokenizer: AutoTokenizer,
        max_batch_size: int = 16,
        max_seq_len: Optional[int] = None,
        max_prompt_len: int = 2048,
        page_size: int = 64,
        device: Optional[str] = None,
        dtype: Optional[torch.dtype] = None,
        cache: Optional[PagePool] = 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:
+        elif config.max_position_embeddings is not None:
-            self.max_seq_len = config.max_len
+            self.max_seq_len = config.max_position_embeddings
        else:
            raise ValueError(
                "max_seq_len must be provided either as argument "
-                "or in model config (config.max_len)"
+                "or in model config (config.max_position_embeddings)"
            )
        self.device = device or next(model.parameters()).device
        self.dtype = dtype or next(model.parameters()).dtype
-        n_pages = (
+        head_dim = config.hidden_size // config.num_attention_heads
            max_batch_size * (self.max_seq_len + page_size) + page_size - 1
        ) // page_size
-        self._page_cache = KVCache(
+        if cache is not None:
-            config.n_layers,
+            self._cache = cache
-            n_pages,
+        else:
-            page_size,
+            self._cache = PagePool(
-            config.n_kv_heads,
+                n_layers=config.num_hidden_layers,
-            config.dim // config.n_heads,
+                n_kv_heads=config.num_key_value_heads,
-            self.device,
+                head_dim=head_dim,
-            self.dtype,
+                max_batch_size=max_batch_size,
-        )
+                max_seq_len=self.max_seq_len,
                device=self.device,
                dtype=self.dtype,
            )
        self._task_mgr = TaskManager(
            tokenizer=tokenizer,
            max_batch_size=max_batch_size,
            max_seq_len=self.max_seq_len,
            max_prompt_len=max_prompt_len,
        )
        self._executor = Executor(
            model=model,
            tokenizer=tokenizer,
-            page_cache=self._page_cache,
+            kv_cache=self._cache,
            device=self.device,
            dtype=self.dtype,
        )
-        self._running = False
+        self._stop_event = threading.Event()
        self._loop_thread: Optional[threading.Thread] = None
    def add_task(self, prompt: str, **kwargs) -> str:
        return self._task_mgr.add_task(prompt, **kwargs)
    def remove_task(self, task_id: str):
        for task in self._task_mgr.remove_task(task_id):
-            self._page_cache.task_free(task.task_id)
+            self._cache.task_free(task.task_id)
    def get_stats(self) -> Dict[str, Any]:
        return self._task_mgr.get_stats()
    def _step(
        self, tasks: List[Task], return_logprobs: bool = False
    ) -> Tuple[List[Task], List[Task]]:
        """Advance every active task by one token (prefill + decode).
        Single shared primitive for both the continuous-batching loop and
        the synchronous ``run_batch`` path, so the two cannot drift.
        Tasks must already be allocated in the KV cache. Tasks without output
        are prefilled first and sample their first token from the final prompt
        position. Tasks with output extend the cache by one position and decode
        from their latest generated token.
        Args:
            tasks: Active tasks to advance by one token.
            return_logprobs: Forwarded to ``execute_decode``; per-token
                logprobs are recorded on each task's ``output_logprobs``.
        Returns:
            ``(decoded, aborted)``: tasks that produced a new token (its ID
            already appended to ``output_ids``) and tasks that hit the
            sequence cap and were marked ``ABORTED``.
        """
        cache = self._cache
        to_prefill = [t for t in tasks if t.output_tokens == 0 and t.prompt_ids]
        prefilled_ids = set()
        produced: List[Task] = []
        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:
                start_pos = min(cache.task_cached(t.task_id), len(t.prompt_ids) - 1)
                groups.setdefault((len(t.prompt_ids), start_pos), []).append(t)
            for (prompt_len, start_pos), group in groups.items():
                prefilled, step_out = self._executor.execute_prefill(
                    group, prompt_len, start_pos, return_logprobs=return_logprobs
                )
                for t, out in zip(prefilled, step_out):
                    t.output_ids.append(out[0] if return_logprobs else out)
                    t.output_tokens += 1
                    prefilled_ids.add(t.task_id)
                    produced.append(t)
                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
                    )
        decoded: List[Task] = []
        aborted: List[Task] = []
        for t in tasks:
            if t.task_id in prefilled_ids:
                continue
            if cache.task_extend(t.task_id, t.next_pos):
                decoded.append(t)
            else:
                t.status = TaskStatus.ABORTED
                aborted.append(t)
        if decoded:
            step_out = self._executor.execute_decode(
                decoded, return_logprobs=return_logprobs
            )
            for t, out in zip(decoded, step_out):
                t.output_ids.append(out[0] if return_logprobs else out)
                t.output_tokens += 1
                produced.append(t)
        return produced, aborted
    def _run_generation_loop(self):
        stop_ids = self._task_mgr.tokenizer.stop_ids
        cache = self._cache
        try:
-            while self._running:
+            while not self._stop_event.is_set():
                finished = self._task_mgr.remove_finished_tasks(stop_ids)
                for task in finished:
-                    self._page_cache.task_free(task.task_id)
+                    if task.status == TaskStatus.FINISHED:
                        cache.task_record_hashes(
                            task.task_id,
                            cache.task_cacheable_ids(
                                task.task_id, task.prompt_ids, task.output_ids
                            ),
                        )
                    cache.task_free(task.task_id)
                active = self._task_mgr.get_active_tasks()
                available = self._task_mgr.max_batch_size - len(active)
@@ -96,7 +179,7 @@ class InferenceScheduler:
                    candidates = self._task_mgr.pull_candidates(available)
                    failed = []
                    for task in candidates:
-                        if self._page_cache.task_alloc(task.task_id, task.prompt_ids):
+                        if cache.task_alloc(task.task_id, task.prompt_ids):
                            self._task_mgr.activate(task)
                        else:
                            failed.append(task)
@@ -107,99 +190,140 @@ class InferenceScheduler:
                    self._task_mgr.wait_for_tasks(timeout=1.0)
                    continue
-                to_prefill = [
+                active = self._task_mgr.get_active_tasks()
                    t
                    for t in self._task_mgr.get_active_tasks()
                    if t.output_tokens == 0
                    and self._page_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]] = {}
+                decoded, aborted = self._step(active)
                    for t in to_prefill:
                        key = (
                            len(t.prompt_ids),
                            self._page_cache.task_cached(t.task_id),
                        )
                        groups.setdefault(key, []).append(t)
-                    for (prompt_len, start_pos), group in groups.items():
+                for t in aborted:
-                        self._executor.execute_prefill(group, prompt_len, start_pos)
+                    self._task_mgr.invoke_callback(t.task_id, STOP)
                        start_logical_page = start_pos // self._page_cache.page_size
                        for t in group:
                            self._page_cache.task_record_hashes(
                                t.task_id,
                                t.prompt_ids,
                                start_logical_page=start_logical_page,
                            )
-                pos_groups: Dict[int, List[Task]] = {}
+                for t in decoded:
-                for t in self._task_mgr.get_active_tasks():
+                    new_text = t.decode_new_token(self._task_mgr.tokenizer)
-                    pos_groups.setdefault(t.next_pos, []).append(t)
+                    if new_text:
-
+                        self._task_mgr.invoke_callback(t.task_id, new_text)
-                if pos_groups:
+                    if t.is_finished(stop_ids):
-                    best_key = max(pos_groups, key=lambda k: len(pos_groups[k]))
+                        remaining = t.flush_remaining(self._task_mgr.tokenizer)
-                    group = sorted(pos_groups[best_key], key=lambda t: t.task_id)
+                        if remaining:
-
+                            self._task_mgr.invoke_callback(t.task_id, remaining)
-                    valid: List[Task] = []
+                        self._task_mgr.invoke_callback(t.task_id, STOP)
                    for t in group:
                        if self._page_cache.task_extend(t.task_id, t.next_pos):
                            valid.append(t)
                        else:
                            t.status = TaskStatus.ABORTED
                            if t.stream_callback:
                                t.stream_callback(STOP)
                    if valid:
                        next_tokens = self._executor.execute_decode(valid)
                        for t, ntok in zip(valid, next_tokens):
                            t.output_ids.append(ntok)
                            t.output_tokens += 1
                            pos = t.input_tokens + t.output_tokens
                            extend_ok = self._page_cache.task_extend(t.task_id, pos)
                            if t.stream_callback:
                                t.stream_callback(
                                    self._task_mgr.tokenizer.decode([ntok])
                                )
                            if not extend_ok:
                                t.status = TaskStatus.ABORTED
                                if t.stream_callback:
                                    t.stream_callback(STOP)
                        for t in valid:
                            if t.is_finished(stop_ids):
                                if t.stream_callback:
                                    t.stream_callback(STOP)
        except Exception as e:
            self._stop_event.set()
            logger.error(f"Scheduler loop crashed: {e}", exc_info=True)
            for task in self._task_mgr.get_active_tasks():
-                if task.stream_callback:
+                self._task_mgr.invoke_callback(task.task_id, STOP)
-                    task.stream_callback(STOP)
+                cache.task_free(task.task_id)
                self._page_cache.task_free(task.task_id)
            for task in self._task_mgr.get_waiting_tasks():
-                if task.stream_callback:
+                self._task_mgr.invoke_callback(task.task_id, STOP)
                    task.stream_callback(STOP)
            self._task_mgr.clear_queues()
            raise
    def start(self):
-        if not self._running:
+        if self._loop_thread is not None and self._loop_thread.is_alive():
-            self._running = True
+            return
-            t = threading.Thread(target=self._run_generation_loop, daemon=True)
+        self._stop_event.clear()
-            t.start()
+        t = threading.Thread(target=self._run_generation_loop, daemon=True)
-            self._loop_thread = t
+        t.start()
        self._loop_thread = t
    def stop(self):
-        self._running = False
+        self._stop_event.set()
        self._task_mgr.wake()
-        if hasattr(self, "_loop_thread"):
+        if self._loop_thread is not None:
            self._loop_thread.join(timeout=2.0)
        self._loop_thread = None
        for task in self._task_mgr.get_active_tasks():
-            self._page_cache.task_free(task.task_id)
+            self._task_mgr.invoke_callback(task.task_id, STOP)
            self._cache.task_free(task.task_id)
        for task in self._task_mgr.get_waiting_tasks():
            self._task_mgr.invoke_callback(task.task_id, STOP)
            self._cache.task_free(task.task_id)
        self._task_mgr.clear_queues()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
    def run_batch(
        self,
        prompt_ids_list: List[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,
        return_logprobs: bool = False,
    ) -> List[List[int]]:
        """Synchronous batch generation without the scheduler thread.
        Accepts already-tokenized prompts (no string round-trip) and runs
        prefill + decode to completion on the calling thread.  Designed for
        RL rollout, where logprobs of the behaviour policy must be collected
        alongside generated tokens.
        Args:
            prompt_ids_list: ``B`` prompts, each a list of token IDs.
            max_tokens: Maximum tokens to generate per prompt.  ``None``
                uses ``self.max_seq_len - len(prompt_ids)``.
            temperature/top_p/top_k/frequency_penalty/rep_window: Sampling
                parameters (uniform across the batch).
            return_logprobs: If ``True``, return ``(token_ids, logprobs)``
                tuples per prompt (logprobs aligned 1-to-1 with token_ids).
        Returns:
            ``List[List[int]]`` of generated token IDs per prompt, or —
            when ``return_logprobs`` is ``True`` —
            ``List[Tuple[List[int], List[float]]]``.
        """
        stop_ids = self._task_mgr.tokenizer.stop_ids
        cache = self._cache
        seq_cap = self.max_seq_len
        tasks: List[Task] = []
        for ids in prompt_ids_list:
            if len(ids) >= seq_cap:
                tasks.append(None)
                continue
            t_max = max_tokens
            if t_max is None:
                t_max = seq_cap - len(ids)
            else:
                t_max = min(t_max, seq_cap - len(ids))
            if t_max <= 0:
                tasks.append(None)
                continue
            task = Task(
                task_id=f"batch_{uuid.uuid4().hex[:8]}",
                prompt_ids=list(ids),
                max_tokens=t_max,
                temperature=temperature,
                top_p=top_p,
                top_k=top_k,
                frequency_penalty=frequency_penalty,
                rep_window=rep_window,
            )
            if not cache.task_alloc(task.task_id, task.prompt_ids):
                tasks.append(None)
                continue
            task.input_tokens = len(task.prompt_ids)
            tasks.append(task)
        try:
            live = [t for t in tasks if t is not None]
            while live:
                decoded, _ = self._step(live, return_logprobs=return_logprobs)
                live = [t for t in decoded if not t.is_finished(stop_ids)]
        finally:
            for t in tasks:
                if t is not None:
                    cache.task_free(t.task_id)
        results: List[Any] = []
        for t in tasks:
            if t is None:
                results.append(([], []) if return_logprobs else [])
            elif return_logprobs:
                results.append((list(t.output_ids), list(t.output_logprobs)))
            else:
                results.append(list(t.output_ids))
        return results
@@ -6,6 +6,8 @@ from collections import deque
 from enum import Enum
 from typing import Any, Callable, Deque, Dict, List, Optional
 from tokenizers.decoders import DecodeStream
 from astrai.tokenize.tokenizer import AutoTokenizer
 logger = logging.getLogger(__name__)
@@ -13,6 +15,33 @@ logger = logging.getLogger(__name__)
 STOP = object()
 class StreamDecoder:
    """Incremental decoder backed by the tokenizers library's DecodeStream.
    Delegates to the Rust-native streaming decoder which maintains an
    O(1) bounded token buffer internally (via prefix drain), avoiding
    the O(n²) cost of re-decoding the full history on each step.
    Multi-byte UTF-8 sequences split across token boundaries are
    buffered until complete; ``push`` returns "" while the trailing
    sequence is still incomplete.
    """
    __slots__ = ("_stream", "_tok")
    def __init__(self, tokenizer: AutoTokenizer):
        self._tok = tokenizer._tokenizer
        self._stream = DecodeStream(skip_special_tokens=True)
    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.
        """
        chunk = self._stream.step(self._tok, token_id)
        return chunk or ""
 class TaskStatus(Enum):
    """Task lifecycle states."""
@@ -33,7 +62,8 @@ class Task:
        temperature: float = 1.0,
        top_p: float = 1.0,
        top_k: int = 50,
-        stream_callback: Optional[Callable[[str], None]] = None,
+        frequency_penalty: float = 0.0,
        rep_window: int = 64,
    ):
        self.task_id = task_id
        self.prompt_ids = prompt_ids
@@ -41,18 +71,42 @@ class Task:
        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.output_logprobs: List[float] = []
        self.input_tokens: int = 0
        self.output_tokens: int = 0
        self.arrival_time = time.time()
        self.finish_time: Optional[float] = None
-        self.stream_callback = stream_callback
+        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.
        With the Rust-native DecodeStream, the stream is always in a
        correct state — any completed text was already emitted by the
        last ``push``.  A trailing incomplete multi-byte sequence has no
        valid text to emit, so this is a no-op.
        """
        return ""
    @property
    def next_pos(self) -> int:
-        return self.input_tokens + len(self.output_ids)
+        # The first output is sampled from prefill and enters KV on the next step.
        return self.input_tokens + max(0, len(self.output_ids) - 1)
    def is_finished(self, stop_ids: List[int]) -> bool:
        if self.max_tokens is not None and self.output_tokens >= self.max_tokens:
@@ -70,15 +124,14 @@ class TaskManager:
        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()
@@ -93,14 +146,16 @@ class TaskManager:
        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:
+        if len(prompt_ids) > self.max_seq_len:
-            prompt_ids = prompt_ids[-self.max_prompt_len :]
+            prompt_ids = prompt_ids[-self.max_seq_len :]
-        if len(prompt_ids) >= self.max_seq_len:
+        if len(prompt_ids) > self.max_seq_len:
            if stream_callback:
                stream_callback(STOP)
            return task_id
@@ -117,12 +172,15 @@ class TaskManager:
            temperature=temperature,
            top_p=top_p,
            top_k=top_k,
-            stream_callback=stream_callback,
+            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
@@ -134,8 +192,14 @@ class TaskManager:
                t for t in self.waiting_queue if t.task_id != task_id
            )
            self.active_tasks = [t for t in self.active_tasks if t.task_id != task_id]
            self._callbacks.pop(task_id, None)
        return removed_active
    def invoke_callback(self, task_id: str, token: str):
        cb = self._callbacks.get(task_id)
        if cb:
            cb(token)
    def get_stats(self) -> Dict[str, Any]:
        return {
            "total_tasks": self._total_tasks,
@@ -186,7 +250,10 @@ class TaskManager:
        return bool(self.active_tasks or self.waiting_queue)
    def wait_for_tasks(self, timeout: float = 1.0):
-        self._task_event.clear()
+        with self._lock:
            if self.waiting_queue or self.active_tasks:
                return
            self._task_event.clear()
        self._task_event.wait(timeout=timeout)
    def get_active_tasks(self) -> List[Task]:
@@ -201,6 +268,7 @@ class TaskManager:
        with self._lock:
            self.waiting_queue.clear()
            self.active_tasks.clear()
            self._callbacks.clear()
    def wake(self):
        self._task_event.set()
@@ -0,0 +1,151 @@
 """Pre-allocated buffers for the inference decode hot path.
 Mirrors FlashInfer / SGLang's global workspace pattern: all per-step tensors
 are allocated eagerly at init (nothing is lazy), so the decode step
 reads/writes fixed-address tensors with zero ``torch.empty`` calls during
 the hot loop — a prerequisite for CUDA-graph capture.
 """
 import torch
 from torch import Tensor
 _MAX_SPLITS = 32
 class InferenceWorkspace:
    """Reusable fixed-shape per-step buffers for decode.
    Families of buffers, all sized to ``max_batch_size`` / ``max_seq_len``
    and sliced via views each step:
    - ``decode_mask``: a ``[B, 1, total_len]`` validity mask, the RHS
      ``arange`` pre-computed so only a single ``torch.ge(out=)`` runs per
      step.
    - ``input_ids``: per-step token IDs filled from host (pinned, double-
      buffered so an in-flight async H2D copy never races the next fill).
    - KV-cache bind metadata (``req_pool_indices``, ``seq_lens``,
      ``kv_indptr``, ``inc``, ``out_cache_loc``), written by
      ``PagePool.bind_tasks`` when the Executor passes this workspace.
    - ``decode_o_part`` / ``decode_ml_part``: split-KV partial result buffers
      (mirrors FlashInfer's workspace).  One global alloc, reused by every
      decode step across all layers.  Sliced views are passed to the CUDA
      attention kernel so its internal ``torch.empty`` hot-path alloc goes
      through a stable address (CUDA-graph capturable).
    No re-allocation while the server's bounds are respected.
    """
    def __init__(
        self,
        max_batch_size: int,
        max_seq_len: int,
        max_q_heads: int,
        head_dim: int,
        device: torch.device,
        dtype: torch.dtype,
    ):
        self.max_batch_size = max_batch_size
        self.max_seq_len = max_seq_len
        self.max_q_heads = max_q_heads
        self.head_dim = head_dim
        self.device = device
        self.dtype = dtype
        # ``position_ids[:, None, None] >= arange`` RHS, reused every step.
        self.arange = torch.arange(max_seq_len, device=device)
        # Decode validity mask: [max_batch, 1, max_seq_len] bool.
        self.input_mask = torch.empty(
            (max_batch_size, 1, max_seq_len), dtype=torch.bool, device=device
        )
        # Per-step token IDs.  Values come from host Python lists every
        # step, so the device buffer is pre-allocated (stable address for
        # CUDA-graph capture) and filled via a host staging buffer.  A
        # double buffer keeps a copy in flight from being overwritten by
        # the next fill.
        self.input_ids = torch.empty((max_batch_size,), dtype=torch.long, device=device)
        self._pin = [
            torch.empty((max_batch_size,), dtype=torch.long),
            torch.empty((max_batch_size,), dtype=torch.long),
        ]
        self._pin_idx = 0
        # KV-cache bind metadata (fixed shape, written by ``PagePool.bind_tasks``
        # when the Executor passes this workspace).  Stable addresses make the
        # decode forward CUDA-graph capturable.
        self.req_pool_indices = torch.empty(
            (max_batch_size,), dtype=torch.long, device=device
        )
        self.seq_lens = torch.empty((max_batch_size,), dtype=torch.long, device=device)
        self.kv_indptr = torch.empty(
            (max_batch_size + 1,), dtype=torch.int32, device=device
        )
        self.qo_indptr = torch.empty(
            (max_batch_size + 1,), dtype=torch.int32, device=device
        )
        self.inc = torch.arange(max_batch_size + 1, dtype=torch.int32, device=device)
        self.out_cache_loc = torch.empty(
            (max_batch_size, 1), dtype=torch.long, device=device
        )
        # Per-step position IDs (must be at a fixed address for CUDA-graph capture).
        self.position_ids = torch.empty(
            (max_batch_size,), dtype=torch.long, device=device
        )
        # Split-KV partial-result buffers for decode (persistent, one global
        # alloc per process — mirrors FlashInfer's workspace pattern).
        # Shape: [max_batch_size, max_q_heads, _MAX_SPLITS, head_dim] (o_part)
        #        [max_batch_size, max_q_heads, _MAX_SPLITS, 2]     (ml_part)
        self.decode_o_part = torch.empty(
            (max_batch_size, max_q_heads, _MAX_SPLITS, head_dim),
            dtype=torch.float32,
            device=device,
        )
        self.decode_ml_part = torch.empty(
            (max_batch_size, max_q_heads, _MAX_SPLITS, 2),
            dtype=torch.float32,
            device=device,
        )
        # Decode output buffer (graph-safe pre-alloc).  Shape matches the
        # decode kernel's output: [batch, q_head, head_dim].
        self.decode_out = torch.empty(
            (max_batch_size, max_q_heads, head_dim),
            dtype=dtype,
            device=device,
        )
    def decode_buffers(self, batch: int, q_heads: int):
        """Return ``(o_part, ml_part)`` view sliced to live dimensions."""
        return (
            self.decode_o_part[:batch, :q_heads],
            self.decode_ml_part[:batch, :q_heads],
        )
    def fill_input_ids(self, ids: "list[int]") -> Tensor:
        """Write ``ids`` into the device buffer and return ``[B]``.
        Host values are staged through the double buffer and copied into the
        stable device buffer (``copy_`` without pinning is synchronous, so
        the alternating buffers guard against an in-flight transfer).
        """
        b = len(ids)
        pin = self._pin[self._pin_idx]
        self._pin_idx ^= 1
        for i, v in enumerate(ids):
            pin[i] = v
        self.input_ids[:b].copy_(pin[:b])
        return self.input_ids[:b]
    def decode_mask(self, position_ids: Tensor, total_len: int) -> Tensor:
        """Return the ``[B, 1, total_len]`` validity mask for this step.
        Written into the pre-allocated buffer via ``torch.ge(out=)`` — no
        new tensor is allocated.  ``position_ids`` is the current step's
        ``[B]`` positions; ``total_len`` must not exceed ``max_seq_len``.
        """
        b = position_ids.size(0)
        out = self.input_mask[:b, :, :total_len]
        torch.ge(position_ids[:, None, None], self.arange[:total_len], out=out)
        return out
@@ -8,6 +8,7 @@ from typing import Any, AsyncGenerator, Dict, Generator, List, Optional, Tuple,
 import torch
 import torch.nn as nn
 from astrai.inference.core.cache import PagePool
 from astrai.inference.core.scheduler import InferenceScheduler
 from astrai.inference.core.task import STOP
 from astrai.tokenize import AutoTokenizer
@@ -63,33 +64,6 @@ class GenerateResult:
            return self.results.copy()
 class GenerationRequest:
    """Request parameters for text generation."""
    def __init__(
        self,
        messages: List[Dict[str, str]],
        top_k: int = 50,
        top_p: float = 1.0,
        temperature: float = 1.0,
        max_tokens: Optional[int] = None,
        stream: bool = False,
    ):
        if not (isinstance(top_k, int) and top_k >= 0):
            raise ValueError("top_k must be a non-negative integer")
        if not (0.0 <= top_p <= 1.0):
            raise ValueError("top_p must be a float between 0.0 and 1.0")
        if not (isinstance(temperature, (int, float)) and temperature >= 0):
            raise ValueError("temperature must be a non-negative number")
        self.messages = messages
        self.top_k = top_k
        self.top_p = top_p
        self.temperature = temperature
        self.max_tokens = max_tokens
        self.stream = stream
 class InferenceEngine:
    """Unified inference engine backed by continuous-batching scheduler."""
@@ -99,8 +73,7 @@ class InferenceEngine:
        tokenizer: AutoTokenizer,
        max_batch_size: int = 1,
        max_seq_len: Optional[int] = None,
-        max_prompt_len: int = 2048,
+        cache: Optional[PagePool] = None,
        page_size: int = 128,
    ):
        self.model = model
        self.tokenizer = tokenizer
@@ -109,8 +82,7 @@ class InferenceEngine:
            tokenizer=self.tokenizer,
            max_batch_size=max_batch_size,
            max_seq_len=max_seq_len,
-            max_prompt_len=max_prompt_len,
+            cache=cache,
            page_size=page_size,
        )
        self.scheduler.start()
@@ -130,18 +102,29 @@ class InferenceEngine:
        temperature: float = 1.0,
        top_p: float = 1.0,
        top_k: int = 50,
        frequency_penalty: float = 0.0,
        rep_window: int = 64,
    ) -> Union[Generator, str, List[str]]:
        is_batch = isinstance(prompt, list)
        prompts = prompt if is_batch else [prompt]
-        if stream:
+        if max_tokens is not None and max_tokens <= 0:
-            return self._generate_streaming(
+            if stream:
-                prompts, is_batch, max_tokens, temperature, top_p, top_k
+                return iter(())
-            )
+            results = [""] * len(prompts)
-        else:
+            return results if is_batch else results[0]
-            return self._generate_non_streaming(
+
-                prompts, is_batch, max_tokens, temperature, top_p, top_k
+        return self._generate(
-            )
+            prompts,
            is_batch,
            stream,
            max_tokens,
            temperature,
            top_p,
            top_k,
            frequency_penalty,
            rep_window,
        )
    def generate_async(
        self,
@@ -150,85 +133,72 @@ class InferenceEngine:
        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]:
-        sync_gen = self._generate_streaming(
+        sync_gen = self._generate(
-            [prompt], False, max_tokens, temperature, top_p, top_k
+            [prompt],
            False,
            True,
            max_tokens,
            temperature,
            top_p,
            top_k,
            frequency_penalty,
            rep_window,
        )
        async def _agen():
            loop = asyncio.get_event_loop()
            while True:
-                token = await loop.run_in_executor(None, self._next_token, sync_gen)
+                try:
-                if token is None:
+                    token = await loop.run_in_executor(None, next, sync_gen)
                except StopIteration:
                    break
                yield token
        return _agen()
-    @staticmethod
+    def _generate(
    def _next_token(gen: Generator) -> Optional[str]:
        try:
            return next(gen)
        except StopIteration:
            return None
    def generate_with_request(
        self, request: GenerationRequest
    ) -> Union[Generator[str, None, None], str, List[str]]:
        prompt = self.tokenizer.apply_chat_template(request.messages, tokenize=False)
        return self.generate(
            prompt=prompt,
            stream=request.stream,
            max_tokens=request.max_tokens,
            temperature=request.temperature,
            top_p=request.top_p,
            top_k=request.top_k,
        )
    def _submit_tasks(
        self,
        prompts: List[str],
        is_batch: bool,
        stream: bool,
        max_tokens: Optional[int],
        temperature: float,
        top_p: float,
        top_k: int,
-    ) -> Tuple[GenerateResult, List[str]]:
+        frequency_penalty: float,
        rep_window: int,
    ) -> Union[Generator, str, List[str]]:
        n = len(prompts)
        result = GenerateResult(count=n)
-        task_ids = []
+        task_ids = [
-        for i, p in enumerate(prompts):
+            self.scheduler.add_task(
            cb = self._make_callback(result, i)
            task_id = self.scheduler.add_task(
                prompt=p,
                max_tokens=max_tokens,
                temperature=temperature,
                top_p=top_p,
                top_k=top_k,
-                stream_callback=cb,
+                frequency_penalty=frequency_penalty,
                rep_window=rep_window,
                stream_callback=lambda token, idx=i: result.append(token, idx),
            )
-            task_ids.append(task_id)
+            for i, p in enumerate(prompts)
-        return result, task_ids
+        ]
-    @staticmethod
+        if not stream:
-    def _make_callback(result: GenerateResult, idx: int):
+            try:
-        def cb(token):
+                result.wait_completion()
-            result.append(token, idx)
+            except TimeoutError:
                for tid in task_ids:
                    self.scheduler.remove_task(tid)
                raise
            for tid in task_ids:
                self.scheduler.remove_task(tid)
            res = result.get_results()
            return res if is_batch else res[0]
        return cb
    def _generate_streaming(
        self,
        prompts: List[str],
        is_batch: bool,
        max_tokens: Optional[int],
        temperature: float,
        top_p: float,
        top_k: int,
    ) -> Generator:
        result, task_ids = self._submit_tasks(
            prompts, max_tokens, temperature, top_p, top_k
        )
        n = len(prompts)
        remaining = n
        finished = [False] * n
@@ -252,32 +222,6 @@ class InferenceEngine:
        return gen()
    def _generate_non_streaming(
        self,
        prompts: List[str],
        is_batch: bool,
        max_tokens: Optional[int],
        temperature: float,
        top_p: float,
        top_k: int,
    ) -> Union[str, List[str]]:
        result, task_ids = self._submit_tasks(
            prompts, max_tokens, temperature, top_p, top_k
        )
        try:
            result.wait_completion()
        except TimeoutError:
            for tid in task_ids:
                self.scheduler.remove_task(tid)
            raise
        for tid in task_ids:
            self.scheduler.remove_task(tid)
        res = result.get_results()
        return res if is_batch else res[0]
    def get_stats(self) -> Dict[str, Any]:
        return self.scheduler.get_stats()
@@ -1,15 +1,15 @@
 """Composable sampling strategies for logit transformation.
 Implements the Strategy pattern: each sampling technique
-(temperature, top-k, top-p) is a pluggable strategy that
+(temperature, top-k, top-p, frequency penalty) is a pluggable
-can be composed into a pipeline.
+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
+from typing import List, Optional, Union
 import torch
 from torch import Tensor
@@ -19,16 +19,28 @@ class BaseSamplingStrategy(ABC):
    """Abstract base for a logit transformation strategy."""
    @abstractmethod
-    def apply(self, logits: Tensor, filter_value: float = -float("inf")) -> Tensor:
+    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):
@@ -41,13 +53,21 @@ class TemperatureStrategy(BaseSamplingStrategy):
    def __init__(self, temperature: Union[float, Tensor] = 1.0):
        self.temperature = temperature
-    def apply(self, logits, filter_value=-float("inf")):
+    def apply(
        self,
        logits: Tensor,
        filter_value: float = -float("inf"),
        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.to(logits.device, non_blocking=True).view(-1, 1)
+                logits = logits / t
        elif t != 1.0:
-            logits = logits / t
+            logits = logits / max(t, 1e-8)
        return logits
@@ -61,7 +81,13 @@ class TopKStrategy(BaseSamplingStrategy):
    def __init__(self, top_k: Union[int, Tensor] = 0):
        self.top_k = top_k
-    def apply(self, logits, filter_value=-float("inf")):
+    def apply(
        self,
        logits: Tensor,
        filter_value: float = -float("inf"),
        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)
@@ -98,7 +124,9 @@ class TopPStrategy(BaseSamplingStrategy):
    def __init__(self, top_p: Union[float, Tensor] = 1.0):
        self.top_p = top_p
-    def _apply(self, logits, top_p, filter_value):
+    def _apply(
        self, logits: Tensor, top_p: Union[float, Tensor], filter_value: float
    ) -> Tensor:
        sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
        cum_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
        remove = cum_probs > top_p
@@ -109,7 +137,13 @@ class TopPStrategy(BaseSamplingStrategy):
        logits[mask] = filter_value
        return logits
-    def apply(self, logits, filter_value=-float("inf")):
+    def apply(
        self,
        logits: Tensor,
        filter_value: float = -float("inf"),
        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)
@@ -120,6 +154,84 @@ class TopPStrategy(BaseSamplingStrategy):
        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.
@@ -140,25 +252,76 @@ class SamplingPipeline(BaseSamplingStrategy):
    def __init__(self, strategies: List[BaseSamplingStrategy]):
        self.strategies = strategies
-    def apply(self, logits, filter_value=-float("inf")):
+    def apply(
        self,
        logits: Tensor,
        filter_value: float = -float("inf"),
        input_ids: Optional[Tensor] = None,
        input_mask: Optional[Tensor] = None,
    ) -> Tensor:
        for strategy in self.strategies:
-            logits = strategy.apply(logits, filter_value)
+            logits = strategy.apply(logits, filter_value, input_ids, input_mask)
        return logits
-    @torch.no_grad()
+    @staticmethod
-    def sample(self, logits: Tensor, filter_value: float = -float("inf")) -> Tensor:
+    def _is_greedy(temperature: Union[float, Tensor]) -> bool:
        if isinstance(temperature, Tensor):
            return bool((temperature == 0).all())
        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,
        return_logprobs: bool = False,
    ):
        """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.
            return_logprobs: If ``True``, return ``(tokens, logprobs)``
                where ``logprobs[i]`` is the log-probability of
                ``tokens[i]`` under the (post-strategy) sampling
                distribution.
        Returns:
-            Sampled token IDs ``[batch]``.
+            Sampled token IDs ``[batch]``, or — when ``return_logprobs``
            is ``True`` — a ``(token_ids, chosen_logprobs)`` tuple.
        """
-        return torch.multinomial(
+        if self._is_greedy_pipeline():
-            torch.softmax(self.apply(logits, filter_value), dim=-1),
+            tokens = logits.argmax(dim=-1)
-            num_samples=1,
+            if not return_logprobs:
                return tokens
            log_probs = torch.log_softmax(logits.float(), dim=-1)
            chosen = torch.gather(log_probs, -1, tokens.unsqueeze(-1)).squeeze(-1)
            return tokens, chosen
        transformed = self.apply(logits, filter_value, input_ids, input_mask)
        tokens = torch.multinomial(
            torch.softmax(transformed, dim=-1), num_samples=1
        ).squeeze(-1)
        if not return_logprobs:
            return tokens
        log_probs = torch.log_softmax(transformed.float(), dim=-1)
        chosen = torch.gather(log_probs, -1, tokens.unsqueeze(-1)).squeeze(-1)
        return tokens, chosen
    def _is_greedy_pipeline(self) -> bool:
        """True if the first strategy is greedy temperature (temp=0)."""
        if not self.strategies:
            return False
        first = self.strategies[0]
        return isinstance(first, TemperatureStrategy) and self._is_greedy(
            first.temperature
        )
@torch.inference_mode()
@@ -167,22 +330,71 @@ def sample(
    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:
+    return_logprobs: bool = False,
 ):
    """Apply sampling strategies then sample (softmax + multinomial).
-    Shortcut for ``SamplingPipeline(...).sample(logits)``.
+    Shortcut for ``SamplingPipeline(...).sample(logits, return_logprobs=)``.
    When **temperature** is exactly 0 (scalar or single-element tensor)
    the function short-circuits to ``argmax`` for deterministic decode.
    When **frequency_penalty** is 0 (the common decode case), the entire
    frequency penalty computation — including the O(batch * vocab) count
    tensor allocation — is skipped.
    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.
        return_logprobs: If ``True``, also return the log-probability
            of each sampled token under the (post-strategy) sampling
            distribution — useful for RL rollout (PPO/GRPO importance
            ratios).
    Returns:
-        Sampled token IDs ``[batch]``.
+        Sampled token IDs ``[batch]``, or — when ``return_logprobs`` is
        ``True`` — a ``(token_ids, chosen_logprobs)`` tuple where
        ``chosen_logprobs`` has shape ``[batch]``.
    """
-    return SamplingPipeline(
+    greedy = (
-        [
+        bool((temperature == 0).all())
-            TemperatureStrategy(temperature),
+        if isinstance(temperature, Tensor)
-            TopKStrategy(top_k),
+        else temperature == 0
-            TopPStrategy(top_p),
+    )
-        ]
+
-    ).sample(logits, filter_value)
+    if greedy:
        tokens = logits.argmax(dim=-1)
        if not return_logprobs:
            return tokens
        log_probs = torch.log_softmax(logits.float(), dim=-1)
        chosen = torch.gather(log_probs, -1, tokens.unsqueeze(-1)).squeeze(-1)
        return tokens, chosen
    has_freq = (
        (isinstance(frequency_penalty, Tensor) and (frequency_penalty != 0).any())
        if isinstance(frequency_penalty, Tensor)
        else frequency_penalty != 0
    )
    strategies: List[BaseSamplingStrategy] = [
        TemperatureStrategy(temperature),
        TopKStrategy(top_k),
        TopPStrategy(top_p),
    ]
    if has_freq:
        strategies.append(FrequencyPenaltyStrategy(frequency_penalty))
    return SamplingPipeline(strategies).sample(
        logits,
        filter_value=filter_value,
        input_ids=input_ids,
        input_mask=input_mask,
        return_logprobs=return_logprobs,
    )
@@ -9,7 +9,7 @@ from astrai.model.components.lora import (
    merge_lora,
    save_lora,
 )
-from astrai.model.components.mlp import MLP
+from astrai.model.components.mlp import MLP, DeepSeekMoE
 from astrai.model.components.norm import RMSNorm
 from astrai.model.encoder import EmbeddingEncoder
 from astrai.model.transformer import AutoRegressiveLM
@@ -19,6 +19,7 @@ __all__ = [
    "Linear",
    "RMSNorm",
    "MLP",
    "DeepSeekMoE",
    "GQA",
    "DecoderBlock",
    # Models
@@ -40,11 +40,12 @@ def _disable_random_init(enable: bool = True):
            setattr(nn.init, n, fn)
-class AutoModel(BaseFactory["AutoModel"], nn.Module):
+class ModelFactory(BaseFactory[nn.Module]):
-    """
+    """Pure factory for model dispatch, separated from nn.Module state."""
-    Autoregressive language model base class.
+
-    Provides model loading/saving, registration, and generation.
+
-    """
+class AutoModel(nn.Module):
    """Model base class with loading/saving and generation."""
    def __init__(self, config: BaseModelConfig):
        super().__init__()
@@ -68,7 +69,7 @@ class AutoModel(BaseFactory["AutoModel"], nn.Module):
        config = ConfigFactory.load(raw)
        model_type = config.model_type or "autoregressive_lm"
-        actual_cls = AutoModel.get_component_class(model_type)
+        actual_cls = ModelFactory.get_component_class(model_type)
        with _disable_random_init(enable=disable_random_init):
            model = actual_cls(config)
@@ -1,12 +1,12 @@
-from astrai.model.components.attention import GQA, MLA, repeat_kv
+from astrai.extension.rotary_backend import apply_rotary_emb
 from astrai.model.components.attention import GQA, MLA
 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.mlp import MLP, DeepSeekMoE
 from astrai.model.components.norm import RMSNorm
 from astrai.model.components.rope import (
    RotaryEmbedding,
    apply_rotary_emb,
    get_rotary_emb,
 )
@@ -14,6 +14,7 @@ __all__ = [
    "Linear",
    "RMSNorm",
    "MLP",
    "DeepSeekMoE",
    "Embedding",
    "GQA",
    "MLA",
@@ -21,5 +22,4 @@ __all__ = [
    "RotaryEmbedding",
    "apply_rotary_emb",
    "get_rotary_emb",
    "repeat_kv",
 ]
@@ -5,28 +5,16 @@ import torch.nn as nn
 import torch.nn.functional as F
 from torch import Tensor
 from astrai.extension import attention
 from astrai.extension.rotary_backend import apply_rotary_emb
 from astrai.factory import BaseFactory
-from astrai.inference.core.cache import KvcacheView
+from astrai.inference.core.cache import KVCache
 from astrai.model.components.linear import Linear
 from astrai.model.components.norm import RMSNorm
 from astrai.model.components.rope import apply_rotary_emb
 def repeat_kv(x: Tensor, n_rep: int) -> Tensor:
    bs, slen, n_heads, head_dim = x.shape
    if n_rep == 1:
        return x
    return (
        x[:, :, :, None, :]
        .expand(bs, slen, n_heads, n_rep, head_dim)
        .reshape(bs, slen, n_heads * n_rep, head_dim)
    )
 class AttnFactory(BaseFactory[nn.Module]):
-    @classmethod
+    pass
    def create(cls, attn_type: str, **kwargs) -> nn.Module:
        return super().create(attn_type, **kwargs)
@AttnFactory.register("gqa")
@@ -40,6 +28,7 @@ class GQA(nn.Module):
        norm_eps: float,
        use_gated_attention: bool,
        layer_id: int,
        n_layers: int = 1,
    ):
        super().__init__()
        assert dim % n_heads == 0
@@ -57,7 +46,7 @@ class GQA(nn.Module):
        self.q_proj = Linear(dim, n_heads * self.head_dim)
        self.k_proj = Linear(dim, n_kv_heads * self.head_dim)
        self.v_proj = Linear(dim, n_kv_heads * self.head_dim)
-        self.o_proj = Linear(dim, dim)
+        self.o_proj = Linear(dim, dim, init_std=0.02 / (2 * n_layers) ** 0.5)
        if self.use_qk_norm:
            self.q_norm = RMSNorm(self.head_dim, norm_eps)
@@ -76,10 +65,9 @@ class GQA(nn.Module):
        x: Tensor,
        rotary_emb: Tensor,
        attn_mask: Tensor = None,
-        paged_cache: Optional[KvcacheView] = None,
+        kv_cache: Optional[KVCache] = None,
        is_causal: bool = False,
    ) -> 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)
@@ -88,19 +76,7 @@ class GQA(nn.Module):
        if self.use_qk_norm:
            q, k = self.q_norm(q), self.k_norm(k)
-        if paged_cache is not None:
+        sdqa_out = attention(q, k, v, kv_cache, self.layer_id, attn_mask, is_causal)
            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))
@@ -123,6 +99,7 @@ class MLA(nn.Module):
        use_qk_norm: bool,
        use_gated_attention: bool,
        layer_id: int,
        n_layers: int = 1,
    ):
        super().__init__()
        self.dim = dim
@@ -150,7 +127,9 @@ class MLA(nn.Module):
            n_kv_heads * (2 * self.head_dim),
        )
-        self.o_proj = Linear(dim, dim, bias=False)
+        self.o_proj = Linear(
            dim, dim, bias=False, init_std=0.02 / (2 * n_layers) ** 0.5
        )
        if use_gated_attention:
            self.gate = Linear(dim, dim, bias=False)
@@ -160,10 +139,10 @@ class MLA(nn.Module):
        x: Tensor,
        rotary_emb: Tensor,
        attn_mask: Tensor = None,
-        paged_cache: Optional[KvcacheView] = None,
+        kv_cache: Optional[KVCache] = None,
        is_causal: bool = False,
    ) -> 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)
@@ -192,18 +171,7 @@ class MLA(nn.Module):
            q = self.q_norm(q)
            k = self.k_norm(k)
-        if paged_cache is not None:
+        attn_out = attention(q, k, v, kv_cache, self.layer_id, attn_mask, is_causal)
            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))
@@ -1,59 +1,74 @@
-from typing import Optional
+from dataclasses import asdict
 from typing import Optional, TypedDict
 import torch.nn as nn
 from torch import Tensor
-from astrai.inference.core.cache import KvcacheView
+from astrai.inference.core.cache import KVCache
 from astrai.model.components.attention import AttnFactory
-from astrai.model.components.mlp import FFNFactory
+from astrai.model.components.mlp import FFNFactory, RouterStats
 from astrai.model.components.norm import RMSNorm
 class DecoderOutput(TypedDict):
    hidden_states: Tensor
    aux_loss: Optional[Tensor]
    router_stats: Optional[RouterStats]
 class DecoderBlock(nn.Module):
-    def __init__(
+    def __init__(self, config, layer_id: int):
        self,
        dim: int,
        n_heads: int,
        dim_ffn: int,
        n_kv_heads: int,
        norm_eps: float,
        use_qk_norm: bool,
        use_gated_attention: bool,
        layer_id: int,
        attn_type: str = "gqa",
        ffn_type: str = "mlp",
        **kwargs,
    ):
        super().__init__()
-        self.attention = AttnFactory.create(
+        cfg = asdict(config)
-            attn_type,
+        cfg.update(
-            dim=dim,
+            dim=config.hidden_size,
-            n_heads=n_heads,
+            dim_ffn=config.intermediate_size,
-            n_kv_heads=n_kv_heads,
+            n_layers=config.num_hidden_layers,
-            use_qk_norm=use_qk_norm,
+            n_heads=config.num_attention_heads,
-            norm_eps=norm_eps,
+            n_kv_heads=config.num_key_value_heads,
-            use_gated_attention=use_gated_attention,
+            norm_eps=config.rms_norm_eps,
-            layer_id=layer_id,
+            down_init_std=0.02 / (2 * config.num_hidden_layers) ** 0.5,
            **kwargs,
        )
-        self.input_norm = RMSNorm(dim, norm_eps)
+        self.attention = AttnFactory.create(config.attn_type, **cfg, layer_id=layer_id)
-        self.post_attention_norm = RMSNorm(dim, norm_eps)
+        self.input_norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
-        self.mlp = FFNFactory.create(ffn_type, dim, dim_ffn, **kwargs)
+        self.post_attention_norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
        ffn_type = self._resolve_ffn_type(config, layer_id)
        self.mlp = FFNFactory.create(ffn_type, **cfg)
    @staticmethod
    def _resolve_ffn_type(config, layer_id: int) -> str:
        if config.ffn_type != "moe":
            return config.ffn_type
        mlp_only = config.mlp_only_layers or []
        if layer_id in mlp_only:
            return "mlp"
        if config.decoder_sparse_step > 1:
            if (layer_id + 1) % config.decoder_sparse_step != 0:
                return "mlp"
        return "moe"
    def forward(
        self,
        x: Tensor,
        rotary_emb: Tensor,
        attention_mask: Optional[Tensor] = None,
-        paged_cache: Optional[KvcacheView] = None,
+        kv_cache: Optional[KVCache] = None,
-    ) -> Tensor:
+        is_causal: bool = False,
    ) -> DecoderOutput:
        attn_output = self.attention(
            self.input_norm(x),
            rotary_emb,
            attention_mask,
-            paged_cache,
+            kv_cache,
            is_causal,
        )
        x = attn_output + x
-        x = self.mlp(self.post_attention_norm(x)) + x
+        normalized = self.post_attention_norm(x)
        mlp_output = self.mlp(normalized)
        x = mlp_output["hidden_states"] + x
-        return x
+        return {
            "hidden_states": x,
            "aux_loss": mlp_output["aux_loss"],
            "router_stats": mlp_output.get("router_stats"),
        }
@@ -1,3 +1,5 @@
 import math
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -5,12 +7,20 @@ from torch import Tensor
 class Embedding(nn.Module):
-    def __init__(self, vocab_size: int, embedding_dim: int):
+    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:
-        return F.embedding(x, self.weight)
+        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
@@ -5,13 +5,16 @@ from torch import Tensor
 class Linear(nn.Module):
-    def __init__(self, in_dim: int, out_dim: int, bias: bool = False):
+    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.kaiming_uniform_(self.weight, a=5**0.5)
+        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)
@@ -1,11 +1,12 @@
 import logging
-from dataclasses import asdict, dataclass
+from dataclasses import asdict
 from pathlib import Path
 from typing import Optional, Set
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from pydantic.dataclasses import dataclass
 from astrai.model.components.linear import Linear
 from astrai.serialization import (
@@ -39,8 +40,12 @@ class LoRALinear(nn.Module):
        self.r = r
        self.scaling = alpha / r
-        self.lora_A = nn.Parameter(torch.randn(r, self.weight.shape[1]) / r)
+        device = self.weight.device
-        self.lora_B = nn.Parameter(torch.zeros(self.weight.shape[0], r))
+        dtype = self.weight.dtype
        lora_a = torch.randn(r, self.weight.shape[1], device=device, dtype=dtype) / r
        lora_b = torch.zeros(self.weight.shape[0], r, device=device, dtype=dtype)
        self.lora_A = nn.Parameter(lora_a)
        self.lora_B = nn.Parameter(lora_b)
        self._merged = False
    def forward(self, x):
@@ -1,3 +1,5 @@
 from typing import Optional, TypedDict
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -8,23 +10,43 @@ from astrai.model.components.linear import Linear
 class FFNFactory(BaseFactory[nn.Module]):
-    @classmethod
+    pass
-    def create(cls, ffn_type: str, dim: int, dim_ffn: int, **kwargs) -> nn.Module:
+
-        return super().create(ffn_type, dim, dim_ffn, **kwargs)
+
 class RouterStats(TypedDict):
    """Per-layer MoE routing statistics for training diagnostics.
    Both tensors are detached monitoring data produced during forward.
    """
    probs: Tensor
    topk_indices: Tensor
 class FFNOutput(TypedDict):
    hidden_states: Tensor
    aux_loss: Optional[Tensor]
    router_stats: Optional[RouterStats]
 class RoutedOutput(TypedDict):
    hidden_states: Tensor
    aux_loss: Optional[Tensor]
    router_stats: Optional[RouterStats]
@FFNFactory.register("mlp")
 class MLP(nn.Module):
-    def __init__(self, dim: int, dim_ffn: int):
+    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)
+        self.down = Linear(dim_ffn, dim, init_std=down_init_std)
-    def forward(self, x: Tensor) -> Tensor:
+    def forward(self, x: Tensor) -> FFNOutput:
        gated = self.up(x) * F.silu(self.gate(x))
        out = self.down(gated)
-        return out
+        return {"hidden_states": out, "aux_loss": None, "router_stats": None}
@FFNFactory.register("moe")
@@ -37,6 +59,10 @@ class DeepSeekMoE(nn.Module):
        n_shared_experts: int = 1,
        n_activated_experts: int = 2,
        topk_method: str = "greedy",
        n_layers: int = 1,
        moe_intermediate_size: Optional[int] = None,
        shared_expert_intermediate_size: Optional[int] = None,
        norm_topk_prob: bool = True,
    ):
        super().__init__()
        self.dim = dim
@@ -44,50 +70,108 @@ class DeepSeekMoE(nn.Module):
        self.n_shared_experts = n_shared_experts
        self.n_activated_experts = n_activated_experts
        self.topk_method = topk_method
        self.norm_topk_prob = norm_topk_prob
        expert_dim_ffn = (
            moe_intermediate_size if moe_intermediate_size is not None else dim_ffn
        )
        shared_dim_ffn = (
            shared_expert_intermediate_size
            if shared_expert_intermediate_size is not None
            else dim_ffn
        )
        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) for _ in range(n_shared_experts)]
+            [
                MLP(dim, shared_dim_ffn, down_init_std=down_init_std)
                for _ in range(n_shared_experts)
            ]
        )
        self.routed_experts = nn.ModuleList(
-            [MLP(dim, dim_ffn) for _ in range(n_routed_experts)]
+            [
                MLP(dim, expert_dim_ffn, down_init_std=down_init_std)
                for _ in range(n_routed_experts)
            ]
        )
-    def forward(self, x: Tensor) -> Tensor:
+    def forward(self, x: Tensor) -> FFNOutput:
        include_aux_loss = self.training and torch.is_grad_enabled()
        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)
+        routed_output = self._routed_forward(x_flat, include_aux_loss)
-        out = (shared_out + routed_out).view(bsz, seq_len, dim)
+        out = (shared_out + routed_output["hidden_states"]).view(bsz, seq_len, dim)
-        return out
+        return {
            "hidden_states": out,
            "aux_loss": routed_output["aux_loss"],
            "router_stats": routed_output["router_stats"],
        }
    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
+        return (
            sum(e(x)["hidden_states"] for e in self.shared_experts)
            / self.n_shared_experts
        )
-    def _routed_forward(self, x: Tensor) -> Tensor:
+    def _routed_forward(self, x: Tensor, include_aux_loss: bool) -> RoutedOutput:
        N, D = x.shape
        K = self.n_activated_experts
        E = self.n_routed_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_indices = torch.topk(router_probs, K, dim=-1, sorted=False)
-        topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
+        if self.norm_topk_prob:
            topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
        aux_loss = None
        router_stats = None
        if include_aux_loss:
            expert_load = F.one_hot(topk_indices, num_classes=E).float()
            expert_load = expert_load.mean(dim=(0, 1))
            router_prob = router_probs.float().mean(dim=0)
            aux_loss = E * (expert_load * router_prob).sum()
            router_stats = {
                "probs": router_probs.detach(),
                "topk_indices": topk_indices,
            }
        # Grouped dispatch: sort (token, slot) pairs by expert so each expert
        # consumes one contiguous slice instead of a per-expert mask scan.
        flat_experts = topk_indices.reshape(-1)
        sorted_experts, order = torch.sort(flat_experts)
        flat_tokens = x.repeat_interleave(K, dim=0)[order]
        flat_weights = topk_weights.reshape(-1, 1)[order]
        boundaries = torch.cumsum(
            torch.bincount(sorted_experts, minlength=E), dim=0
        ).tolist()
        output = torch.zeros(N, D, device=x.device, dtype=x.dtype)
-        for expert_idx in range(self.n_routed_experts):
+        start = 0
-            expert_mask = topk_indices == expert_idx
+        for expert_idx, end in enumerate(boundaries):
-            token_idx, k_idx = expert_mask.nonzero(as_tuple=True)
+            if end == start:
            if token_idx.numel() == 0:
                continue
-            expert_input = x[token_idx]
+            expert_output = self.routed_experts[expert_idx](flat_tokens[start:end])[
-            expert_output = self.routed_experts[expert_idx](expert_input)
+                "hidden_states"
-            weights = topk_weights[token_idx, k_idx].unsqueeze(-1)
+            ]
-            output.index_add_(0, token_idx, expert_output * weights)
+            output.index_add_(
                0,
                order[start:end] // K,
                expert_output * flat_weights[start:end],
            )
            start = end
-        return output
+        return {
            "hidden_states": output,
            "aux_loss": aux_loss,
            "router_stats": router_stats,
        }
@@ -11,28 +11,23 @@ def get_rotary_emb(
    base: float = 10000,
    device: Optional[torch.device] = None,
 ) -> Tensor:
    """Precompute cos/sin tables for rotary embedding.
    Returns:
        [max_len, dim/2, 2] (f32) — [cos, sin] pairs.
    """
    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)
+    return torch.stack([cos, sin], dim=-1)
 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,
@@ -56,16 +51,23 @@ class RotaryEmbedding(nn.Module):
        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 = 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:
        """Lookup cos/sin for the given positions.
        Args:
            x: [batch, seq_len, ...] — only batch and seq_len are used.
            position_ids: [batch, seq_len] optional position indices.
        Returns:
            [batch, seq_len, dim/2, 2] (f32) — [cos, sin] pairs.
        """
        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 self.freqs_cis[position_ids].float()
        return torch.view_as_complex(position_freq_cis)
@@ -5,7 +5,7 @@ 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.automodel import AutoModel, ModelFactory
 from astrai.model.components.decoder_block import DecoderBlock
 from astrai.model.components.embedding import Embedding
 from astrai.model.components.norm import RMSNorm
@@ -13,35 +13,33 @@ from astrai.model.components.rope import RotaryEmbedding
 from astrai.model.transformer import process_attention_mask
-@AutoModel.register("embedding")
+@ModelFactory.register("embedding")
 class EmbeddingEncoder(AutoModel):
    def __init__(self, config: EncoderConfig):
        super().__init__(config)
        self.config = config
-        rope_dim = config.dim // config.n_heads
+        rope_dim = config.hidden_size // config.num_attention_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
+            rope_dim,
            config.max_position_embeddings,
            rope_base,
            rope_scaling=config.rope_scaling,
        )
        self.embed_tokens = Embedding(
            config.vocab_size,
            config.hidden_size,
            neftune_alpha=config.neftune_alpha,
        )
        self.embed_tokens = Embedding(config.vocab_size, config.dim)
        self.layers = nn.ModuleList(
            [
-                DecoderBlock(
+                DecoderBlock(config, layer_id)
-                    config.dim,
+                for layer_id in range(config.num_hidden_layers)
                    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.norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
        self.pooling_type = config.pooling_type or "mean"
        self.normalize_embeddings = config.normalize_embeddings or False
@@ -69,10 +67,10 @@ class EmbeddingEncoder(AutoModel):
        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)
+        attn_mask = process_attention_mask(input_mask)
        for layer in self.layers:
-            x = layer(x, rotary_emb, attn_mask, paged_cache=None)
+            x = layer(x, rotary_emb, attn_mask)["hidden_states"]
        hidden_states = self.norm(x)
@@ -5,8 +5,8 @@ import torch.nn as nn
 from torch import Tensor
 from astrai.config.model_config import AutoRegressiveLMConfig
-from astrai.inference.core.cache import KvcacheView
+from astrai.inference.core.cache import KVCache
-from astrai.model.automodel import AutoModel
+from astrai.model.automodel import AutoModel, ModelFactory
 from astrai.model.components.decoder_block import DecoderBlock
 from astrai.model.components.embedding import Embedding
 from astrai.model.components.linear import Linear
@@ -15,38 +15,18 @@ from astrai.model.components.rope import RotaryEmbedding
 def process_attention_mask(
-    input_tensor: Tensor,
+    input_mask: Optional[Tensor],
    position_ids: Optional[Tensor],
    input_mask: Optional[Tensor] = None,
    is_causal: bool = False,
 ) -> Optional[Tensor]:
    if position_ids is None:
        return None
    if input_mask is not None and input_mask.dim() > 2:
        return input_mask
    device = input_tensor.device
    dtype = input_tensor.dtype
    B, S = input_tensor.size()[:2]
    T = position_ids.max().item() + 1
    if input_mask is None:
-        if position_ids.min().item() == 0 and is_causal:
+        return None
-            return None
+    if input_mask.dim() == 2:
-        pad = torch.ones(B, T, dtype=torch.bool, device=device)
+        return input_mask[:, None, None, :]
-    else:
+    if input_mask.dim() == 3:
-        pad = input_mask[:, :T].to(device=device, dtype=torch.bool)
+        return input_mask[:, None, :, :]
-
+    return input_mask
    attend = pad.view(B, 1, T).expand(B, S, T).clone()
    if is_causal:
        attend &= position_ids.unsqueeze(-1) >= torch.arange(T, device=device)
    return torch.full(
        (B, 1, S, T), -torch.finfo(dtype).max / 2, dtype=dtype, device=device
    ).masked_fill_(attend.unsqueeze(1), 0.0)
-@AutoModel.register("autoregressive_lm")
+@ModelFactory.register("autoregressive_lm")
 class AutoRegressiveLM(AutoModel):
    """Autoregressive language model with paged KV cache."""
@@ -56,43 +36,32 @@ class AutoRegressiveLM(AutoModel):
        rope_dim = (
            config.qk_rope_head_dim
            if config.attn_type == "mla"
-            else config.dim // config.n_heads
+            else config.hidden_size // config.num_attention_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
+            rope_dim,
            config.max_position_embeddings,
            rope_base,
            rope_scaling=config.rope_scaling,
        )
        self.embed_tokens = Embedding(
            config.vocab_size,
            config.hidden_size,
            neftune_alpha=config.neftune_alpha,
        )
        self.embed_tokens = Embedding(config.vocab_size, config.dim)
        self.layers = nn.ModuleList(
            [
-                DecoderBlock(
+                DecoderBlock(config, layer_id)
-                    config.dim,
+                for layer_id in range(config.num_hidden_layers)
                    config.n_heads,
                    config.dim_ffn,
                    config.n_kv_heads,
                    config.norm_eps,
                    config.use_qk_norm,
                    config.use_gated_attention,
                    layer_id,
                    attn_type=config.attn_type,
                    ffn_type=config.ffn_type,
                    n_routed_experts=config.n_routed_experts,
                    n_shared_experts=config.n_shared_experts,
                    n_activated_experts=config.n_activated_experts,
                    topk_method=config.topk_method,
                    kv_lora_rank=config.kv_lora_rank,
                    qk_nope_head_dim=config.qk_nope_head_dim,
                    qk_rope_head_dim=config.qk_rope_head_dim,
                )
                for layer_id in range(config.n_layers)
            ]
        )
-        self.norm = RMSNorm(config.dim, config.norm_eps)
+        self.norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
-        self.lm_head = Linear(config.dim, config.vocab_size)
+        self.lm_head = Linear(config.hidden_size, config.vocab_size)
-        if self.config.tie_weight is True:
+        if self.config.tie_word_embeddings is True:
            self.lm_head.weight = self.embed_tokens.weight
        self.apply(self._init_weights)
@@ -107,7 +76,7 @@ class AutoRegressiveLM(AutoModel):
        state_dict = dict(state_dict)
-        if self.config.tie_weight is True:
+        if self.config.tie_word_embeddings is True:
            # same tensor for embed and lm_head
            if embed_key in state_dict:
                state_dict[lm_head_key] = state_dict[embed_key]
@@ -123,7 +92,7 @@ class AutoRegressiveLM(AutoModel):
            destination=destination, prefix=prefix, keep_vars=keep_vars
        )
-        if self.config.tie_weight is True:
+        if self.config.tie_word_embeddings is True:
            lm_head_key = prefix + "lm_head.weight"
            if lm_head_key in state_dict:
                del state_dict[lm_head_key]
@@ -134,19 +103,37 @@ class AutoRegressiveLM(AutoModel):
        self,
        input_ids: Tensor,
        input_mask: Optional[Tensor] = None,
-        paged_cache: Optional[KvcacheView] = None,
+        kv_cache: Optional[KVCache] = None,
        position_ids: Optional[Tensor] = None,
    ) -> Dict[str, Tensor]:
        assert input_ids.ndim == 2
        x = self.embed_tokens(input_ids)
        rotary_emb = self.rotary_embedding(x, position_ids)
-        attn_mask = process_attention_mask(x, position_ids, input_mask, is_causal=True)
+        attn_mask = process_attention_mask(input_mask)
        use_sdpa_causal_mask = attn_mask is None
        aux_losses = []
        router_stats_list = []
        for layer in self.layers:
-            x = layer(x, rotary_emb, attn_mask, paged_cache)
+            layer_output = layer(
                x,
                rotary_emb,
                attn_mask,
                kv_cache,
                use_sdpa_causal_mask,
            )
            x = layer_output["hidden_states"]
            stats = layer_output.get("router_stats")
            if stats is not None:
                aux_losses.append(layer_output["aux_loss"])
                router_stats_list.append(stats)
        hidden_states = self.norm(x)
        logits = self.lm_head(hidden_states)
-        return {"logits": logits, "hidden_states": hidden_states}
+        output = {"logits": logits, "hidden_states": hidden_states}
        if aux_losses:
            output["aux_loss"] = torch.stack(aux_losses).mean()
            output["router_stats"] = router_stats_list
        return output
@@ -0,0 +1,38 @@
 """Optimizer implementations and factory registration."""
 from astrai.optim.composite import (
    OptimizerFactory,
    composite_state_dict,
    composite_step,
    composite_zero_grad,
    refresh_param_groups,
 )
 from astrai.optim.mano_adamw import Mano, ManoAdamW
 from astrai.optim.muon_adamw import MuonAdamW
 from astrai.optim.nora_nadamw import (
    NAdamW,
    Nora,
    NoraNAdamW,
    OptimizerParameterGroups,
    nora_direction,
    nora_lr_scale,
    partition_optimizer_parameters,
 )
 __all__ = [
    "Mano",
    "ManoAdamW",
    "MuonAdamW",
    "NAdamW",
    "Nora",
    "NoraNAdamW",
    "OptimizerFactory",
    "OptimizerParameterGroups",
    "composite_state_dict",
    "composite_step",
    "composite_zero_grad",
    "nora_direction",
    "nora_lr_scale",
    "partition_optimizer_parameters",
    "refresh_param_groups",
 ]
@@ -0,0 +1,71 @@
 """Shared infrastructure for the optim package.
 This module hosts two things:
 * ``OptimizerFactory`` — the registry for built-in optimizers. Defining it
  here (rather than in ``__init__.py``) lets each optimizer module import it
  and register itself with a decorator, avoiding circular imports.
 * Composite-optimizer helpers — ``step``/``zero_grad``/``state_dict``/
  ``param_groups`` delegation shared by every optimizer that routes different
  parameter groups through distinct sub-optimizers.
 """
 from typing import Any
 import torch
 from torch.optim import Optimizer
 from astrai.factory import BaseFactory
 class OptimizerFactory(BaseFactory[Optimizer]):
    """Factory for built-in training optimizers."""
 def composite_step(
    sub_optimizers: list[Optimizer],
    closure=None,
 ) -> torch.Tensor | None:
    """Run ``step`` on every sub-optimizer, invoking the closure once.
    The closure (if given) is executed inside ``torch.enable_grad`` exactly
    once before any sub-optimizer steps, matching the contract of a single
    ``Optimizer.step``. Sub-optimizers receive ``None`` so they do not
    re-execute it.
    """
    loss = None
    if closure is not None:
        with torch.enable_grad():
            loss = closure()
    for sub in sub_optimizers:
        sub.step()
    return loss
 def composite_zero_grad(
    sub_optimizers: list[Optimizer],
    set_to_none: bool = True,
 ) -> None:
    for sub in sub_optimizers:
        sub.zero_grad(set_to_none=set_to_none)
 def composite_state_dict(
    named_sub_optimizers: dict[str, Optimizer | None],
 ) -> dict[str, Any]:
    """Serialize sub-optimizers, preserving ``None`` slots."""
    return {
        name: sub.state_dict() if sub is not None else None
        for name, sub in named_sub_optimizers.items()
    }
 def refresh_param_groups(
    sub_optimizers: list[Optimizer],
 ) -> list[dict]:
    """Concatenate param_groups from every non-None sub-optimizer."""
    groups: list[dict] = []
    for sub in sub_optimizers:
        if sub is not None:
            groups.extend(sub.param_groups)
    return groups
@@ -0,0 +1,214 @@
 """Mano manifold optimizer combined with AdamW.
 Mano projects the momentum onto the tangent space of the Oblique manifold
 (axis-wise tangent projection) and normalizes it, replacing the expensive
 Newton-Schulz iteration in Muon with a cheaper manifold normalization.
 Reference: https://arxiv.org/abs/2601.23000
 """
 import math
 import torch
 from torch import nn, optim
 from torch.optim import Optimizer
 from astrai.optim.composite import (
    OptimizerFactory,
    composite_state_dict,
    composite_step,
    composite_zero_grad,
    refresh_param_groups,
 )
 from astrai.optim.nora_nadamw import partition_optimizer_parameters
 class Mano(Optimizer):
    """Manifold Normalized Optimizer for two-dimensional matrices.
    Each step alternates the projection axis (dim 0 / dim 1) to restrike the
    manifold along both rows and columns. The tangent momentum is computed
    without normalizing the parameter itself (v2 simplification) and the
    epsilon is added (not clamped) to the norm denominator.
    """
    def __init__(
        self,
        params,
        lr: float = 1e-3,
        weight_decay: float = 0.1,
        momentum: float = 0.95,
        nesterov: bool = True,
        eps: float = 1e-8,
    ):
        if lr < 0:
            raise ValueError(f"Invalid learning rate: {lr}")
        if weight_decay < 0:
            raise ValueError(f"Invalid weight decay: {weight_decay}")
        if not 0 <= momentum <= 1:
            raise ValueError(f"Invalid momentum: {momentum}")
        if eps <= 0:
            raise ValueError(f"Invalid epsilon: {eps}")
        defaults = {
            "lr": lr,
            "weight_decay": weight_decay,
            "momentum": momentum,
            "nesterov": nesterov,
            "eps": eps,
            "steps": 0,
        }
        super().__init__(params, defaults)
        for group in self.param_groups:
            for param in group["params"]:
                if param.ndim != 2:
                    raise ValueError(
                        f"Mano only supports 2D matrices, got shape {tuple(param.shape)}"
                    )
    @torch.no_grad()
    def step(self, closure=None):
        loss = None
        if closure is not None:
            with torch.enable_grad():
                loss = closure()
        for group in self.param_groups:
            lr = group["lr"]
            weight_decay = group["weight_decay"]
            momentum = group["momentum"]
            nesterov = group["nesterov"]
            eps = group["eps"]
            dim = int(group["steps"] % 2)
            for param in group["params"]:
                if param.grad is None:
                    continue
                if param.grad.is_sparse:
                    raise RuntimeError("Mano does not support sparse gradients")
                grad = param.grad
                state = self.state[param]
                momentum_buffer = state.get("momentum_buffer")
                if momentum_buffer is None:
                    momentum_buffer = torch.zeros_like(grad)
                momentum_buffer.mul_(momentum).add_(grad)
                update = (
                    grad.add(momentum_buffer, alpha=momentum)
                    if nesterov
                    else momentum_buffer
                )
                tangent = update - (
                    torch.sum(update * param.data, dim=dim, keepdim=True) * param.data
                )
                direction = tangent / (
                    torch.norm(tangent, p=2, dim=dim, keepdim=True) + eps
                )
                if weight_decay != 0:
                    param.mul_(1 - lr * weight_decay)
                adjusted_lr = lr * 0.2 * math.sqrt(direction.shape[dim])
                param.add_(direction, alpha=-adjusted_lr)
                state["momentum_buffer"] = momentum_buffer
            group["steps"] += 1
        return loss
@OptimizerFactory.register("mano_adamw")
 class ManoAdamW(Optimizer):
    """Mano for internal linear weights and AdamW for remaining parameters."""
    optimizer_name = "mano_adamw"
    def __init__(
        self,
        model: nn.Module,
        lr: float = 3e-4,
        weight_decay: float = 0.1,
        momentum: float = 0.95,
        nesterov: bool = True,
    ):
        groups = partition_optimizer_parameters(model)
        all_params = [
            *groups.nora,
            *groups.nadamw_decay,
            *groups.nadamw_no_decay,
        ]
        if not all_params:
            raise ValueError(
                "Cannot build an optimizer for a model with no trainable parameters"
            )
        super().__init__(all_params, {})
        self.mano = (
            Mano(
                groups.nora,
                lr=lr,
                weight_decay=weight_decay,
                momentum=momentum,
                nesterov=nesterov,
            )
            if groups.nora
            else None
        )
        adamw_groups = []
        if groups.nadamw_decay:
            adamw_groups.append(
                {"params": groups.nadamw_decay, "weight_decay": weight_decay}
            )
        if groups.nadamw_no_decay:
            adamw_groups.append({"params": groups.nadamw_no_decay, "weight_decay": 0.0})
        self.adamw = (
            optim.AdamW(
                adamw_groups,
                lr=lr,
                betas=(0.9, 0.95),
                fused=True,
            )
            if adamw_groups
            else None
        )
        self.param_groups = refresh_param_groups([self.mano, self.adamw])
    @torch.no_grad()
    def step(self, closure=None):
        return composite_step(
            [opt for opt in (self.mano, self.adamw) if opt is not None],
            closure,
        )
    def zero_grad(self, set_to_none: bool = True):
        composite_zero_grad(
            [opt for opt in (self.mano, self.adamw) if opt is not None],
            set_to_none,
        )
    def state_dict(self) -> dict:
        return composite_state_dict({"mano": self.mano, "adamw": self.adamw})
    def load_state_dict(self, state_dict: dict):
        if "muon" in state_dict or "nora" in state_dict:
            raise ValueError(
                "Checkpoint uses a different optimizer; select the matching "
                "--optimizer to resume it"
            )
        if "mano" not in state_dict or "adamw" not in state_dict:
            raise ValueError(
                "Checkpoint optimizer state is not compatible with mano_adamw"
            )
        saved_mano = state_dict["mano"]
        saved_adamw = state_dict["adamw"]
        if (self.mano is None) != (saved_mano is None):
            raise ValueError("Checkpoint Mano parameter groups do not match the model")
        if (self.adamw is None) != (saved_adamw is None):
            raise ValueError("Checkpoint AdamW parameter groups do not match the model")
        if self.mano is not None:
            self.mano.load_state_dict(saved_mano)
        if self.adamw is not None:
            self.adamw.load_state_dict(saved_adamw)
        self.param_groups = refresh_param_groups([self.mano, self.adamw])
@@ -0,0 +1,95 @@
 """Legacy Muon + AdamW combined optimizer."""
 from typing import Any
 import torch
 from torch import Tensor, nn, optim
 from astrai.optim.composite import (
    OptimizerFactory,
    composite_state_dict,
    composite_step,
    composite_zero_grad,
    refresh_param_groups,
 )
@OptimizerFactory.register("muon_adamw")
 class MuonAdamW(optim.Optimizer):
    """Combined Muon (matrix) + AdamW (non-matrix) optimizer."""
    optimizer_name = "muon_adamw"
    def __init__(
        self,
        model: nn.Module,
        lr: float = 3e-4,
        weight_decay: float = 0.1,
        momentum: float = 0.95,
        nesterov: bool = True,
        ns_steps: int = 5,
        adjust_lr_fn: str = "match_rms_adamw",
    ):
        defaults = {
            "lr": lr,
            "weight_decay": weight_decay,
            "momentum": momentum,
            "nesterov": nesterov,
            "ns_steps": ns_steps,
            "adjust_lr_fn": adjust_lr_fn,
        }
        params = [param for param in model.parameters() if param.requires_grad]
        super().__init__(params, defaults)
        matrix_params: list[Tensor] = []
        other_params: list[Tensor] = []
        for name, param in model.named_parameters():
            if not param.requires_grad:
                continue
            if (
                param.dim() >= 2
                and "norm" not in name
                and "bias" not in name
                and "embed" not in name
                and "lm_head" not in name
            ):
                matrix_params.append(param)
            else:
                other_params.append(param)
        self.muon = optim.Muon(
            matrix_params,
            lr=lr,
            weight_decay=weight_decay,
            momentum=momentum,
            nesterov=nesterov,
            ns_steps=ns_steps,
            adjust_lr_fn=adjust_lr_fn,
        )
        self.adamw = optim.AdamW(
            [{"params": other_params, "weight_decay": 0.0}],
            lr=lr,
            betas=(0.9, 0.95),
            fused=True,
        )
        self.param_groups = refresh_param_groups([self.muon, self.adamw])
    @torch.no_grad()
    def step(self, closure=None):
        return composite_step([self.muon, self.adamw], closure)
    def zero_grad(self, set_to_none: bool = True):
        composite_zero_grad([self.muon, self.adamw], set_to_none)
    def state_dict(self) -> dict[str, Any]:
        return composite_state_dict({"muon": self.muon, "adamw": self.adamw})
    def load_state_dict(self, state_dict: dict[str, Any]):
        if "muon" not in state_dict or "adamw" not in state_dict:
            raise ValueError(
                "Checkpoint optimizer state is not compatible with muon_adamw"
            )
        self.muon.load_state_dict(state_dict["muon"])
        self.adamw.load_state_dict(state_dict["adamw"])
        self.param_groups = refresh_param_groups([self.muon, self.adamw])
@@ -0,0 +1,372 @@
 """Nora matrix optimizer combined with Nesterov AdamW."""
 import math
 from dataclasses import dataclass
 from typing import Any
 import torch
 from torch import Tensor, nn
 from torch.distributed.tensor import DTensor, Shard
 from torch.optim import Optimizer
 from astrai.model.components.embedding import Embedding
 from astrai.model.components.linear import Linear
 from astrai.model.components.lora import LoRALinear
 from astrai.model.components.norm import RMSNorm
 from astrai.optim.composite import (
    OptimizerFactory,
    composite_state_dict,
    composite_step,
    composite_zero_grad,
    refresh_param_groups,
 )
 NORA_EPS = 1e-10
 def _row_normalize(tensor: Tensor, eps: float) -> Tensor:
    return tensor / tensor.norm(dim=-1, keepdim=True).clamp(min=eps)
 def nora_direction(update: Tensor, param: Tensor, eps: float = NORA_EPS) -> Tensor:
    """Project an update onto each parameter row's tangent space and normalize."""
    theta_hat = _row_normalize(param.to(torch.float32), eps)
    update_fp32 = update.to(torch.float32)
    radial = (update_fp32 * theta_hat).sum(dim=-1, keepdim=True) * theta_hat
    direction = _row_normalize(update_fp32 - radial, eps)
    return direction.to(update.dtype)
 def nora_lr_scale(lr: float, shape: torch.Size) -> float:
    """Scale Nora's LR for tall ``[d_out, d_in]`` linear weights."""
    return lr * math.sqrt(max(1.0, shape[-2] / shape[-1]))
 def _validate_complete_rows(param: Tensor) -> None:
    if not isinstance(param, DTensor):
        return
    last_dim = param.ndim - 1
    for placement in param.placements:
        if isinstance(placement, Shard) and placement.dim % param.ndim == last_dim:
            raise ValueError(
                "Nora requires complete parameter rows, but this DTensor is sharded "
                "along its last dimension"
            )
 class Nora(Optimizer):
    """Normalized Orthogonal Row Alignment for two-dimensional matrices."""
    def __init__(
        self,
        params,
        lr: float = 5e-3,
        weight_decay: float = 0.0,
        momentum: float = 0.95,
        beta: float = 0.95,
        nesterov: bool = True,
        eps: float = NORA_EPS,
    ):
        if lr < 0:
            raise ValueError(f"Invalid learning rate: {lr}")
        if weight_decay < 0:
            raise ValueError(f"Invalid weight decay: {weight_decay}")
        if not 0 <= momentum <= 1:
            raise ValueError(f"Invalid momentum: {momentum}")
        if not 0 <= beta < 1:
            raise ValueError(f"Invalid beta: {beta}")
        if eps <= 0:
            raise ValueError(f"Invalid epsilon: {eps}")
        defaults = {
            "lr": lr,
            "weight_decay": weight_decay,
            "momentum": momentum,
            "beta": beta,
            "nesterov": nesterov,
            "eps": eps,
        }
        super().__init__(params, defaults)
        for group in self.param_groups:
            for param in group["params"]:
                if param.ndim != 2:
                    raise ValueError(
                        f"Nora only supports 2D matrices, got shape {tuple(param.shape)}"
                    )
                _validate_complete_rows(param)
    @torch.no_grad()
    def step(self, closure=None):
        loss = None
        if closure is not None:
            with torch.enable_grad():
                loss = closure()
        for group in self.param_groups:
            lr = group["lr"]
            weight_decay = group["weight_decay"]
            momentum = group["momentum"]
            beta = group["beta"]
            nesterov = group["nesterov"]
            eps = group["eps"]
            for param in group["params"]:
                if param.grad is None:
                    continue
                if param.grad.is_sparse:
                    raise RuntimeError("Nora does not support sparse gradients")
                grad = param.grad
                state = self.state[param]
                momentum_buffer = state.get("momentum_buffer")
                if momentum_buffer is None:
                    momentum_buffer = torch.zeros_like(grad)
                momentum_buffer.lerp_(grad, 1 - beta)
                update = (
                    grad.lerp(momentum_buffer, momentum)
                    if nesterov
                    else momentum_buffer
                )
                direction = nora_direction(update, param, eps)
                if weight_decay != 0:
                    param.mul_(1 - lr * weight_decay)
                param.add_(direction, alpha=-nora_lr_scale(lr, param.shape))
                state["momentum_buffer"] = momentum_buffer
        return loss
 class NAdamW(Optimizer):
    """AdamW using the reference Nesterov first-moment update."""
    def __init__(
        self,
        params,
        lr: float = 3e-4,
        betas: tuple[float, float] = (0.9, 0.999),
        eps: float = 1e-8,
        weight_decay: float = 0.1,
    ):
        beta1, beta2 = betas
        if lr < 0:
            raise ValueError(f"Invalid learning rate: {lr}")
        if not 0 <= beta1 < 1 or not 0 <= beta2 < 1:
            raise ValueError(f"Invalid betas: {betas}")
        if eps <= 0:
            raise ValueError(f"Invalid epsilon: {eps}")
        if weight_decay < 0:
            raise ValueError(f"Invalid weight decay: {weight_decay}")
        defaults = {
            "lr": lr,
            "betas": betas,
            "eps": eps,
            "weight_decay": weight_decay,
        }
        super().__init__(params, defaults)
    @torch.no_grad()
    def step(self, closure=None):
        loss = None
        if closure is not None:
            with torch.enable_grad():
                loss = closure()
        for group in self.param_groups:
            beta1, beta2 = group["betas"]
            eps = group["eps"]
            lr = group["lr"]
            weight_decay = group["weight_decay"]
            for param in group["params"]:
                if param.grad is None:
                    continue
                if param.grad.is_sparse:
                    raise RuntimeError("NAdamW does not support sparse gradients")
                grad = param.grad
                state = self.state[param]
                if not state:
                    state["step"] = 0
                    state["m"] = torch.zeros_like(param)
                    state["v"] = torch.zeros_like(param)
                state["step"] += 1
                first_moment = state["m"]
                second_moment = state["v"]
                first_moment.mul_(beta1).add_(grad, alpha=1 - beta1)
                second_moment.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
                bias_correction1 = 1 - beta1 ** state["step"]
                bias_correction2 = 1 - beta2 ** state["step"]
                nesterov_moment = (
                    beta1 * first_moment + (1 - beta1) * grad
                ) / bias_correction1
                corrected_second_moment = second_moment / bias_correction2
                if weight_decay != 0:
                    param.mul_(1 - lr * weight_decay)
                param.addcdiv_(
                    nesterov_moment,
                    corrected_second_moment.sqrt().add_(eps),
                    value=-lr,
                )
        return loss
@dataclass
 class OptimizerParameterGroups:
    nora: list[Tensor]
    nadamw_decay: list[Tensor]
    nadamw_no_decay: list[Tensor]
 def partition_optimizer_parameters(model: nn.Module) -> OptimizerParameterGroups:
    """Partition trainable parameters by module role and parameter identity."""
    nora_ids: set[int] = set()
    no_decay_ids: set[int] = set()
    for module_name, module in model.named_modules():
        if isinstance(module, LoRALinear):
            for param in module.parameters(recurse=False):
                if param.requires_grad:
                    no_decay_ids.add(id(param))
            continue
        if isinstance(module, (Embedding, RMSNorm)):
            for param in module.parameters(recurse=False):
                if param.requires_grad:
                    no_decay_ids.add(id(param))
            continue
        if not isinstance(module, Linear):
            continue
        if module.bias is not None and module.bias.requires_grad:
            no_decay_ids.add(id(module.bias))
        if not module.weight.requires_grad:
            continue
        if module_name.rsplit(".", 1)[-1] == "lm_head":
            no_decay_ids.add(id(module.weight))
        elif module.weight.ndim == 2:
            nora_ids.add(id(module.weight))
    nora: list[Tensor] = []
    nadamw_decay: list[Tensor] = []
    nadamw_no_decay: list[Tensor] = []
    seen: set[int] = set()
    for param in model.parameters():
        param_id = id(param)
        if not param.requires_grad or param_id in seen:
            continue
        seen.add(param_id)
        if param_id in no_decay_ids or param.ndim <= 1:
            nadamw_no_decay.append(param)
        elif param_id in nora_ids:
            nora.append(param)
        else:
            nadamw_decay.append(param)
    trainable_ids = {id(param) for param in model.parameters() if param.requires_grad}
    grouped_ids = {id(param) for param in [*nora, *nadamw_decay, *nadamw_no_decay]}
    if grouped_ids != trainable_ids:
        missing = len(trainable_ids - grouped_ids)
        extra = len(grouped_ids - trainable_ids)
        raise RuntimeError(
            f"Optimizer parameter partition is incomplete: missing={missing}, extra={extra}"
        )
    return OptimizerParameterGroups(nora, nadamw_decay, nadamw_no_decay)
@OptimizerFactory.register("nora_nadamw")
 class NoraNAdamW(Optimizer):
    """Nora for internal linear weights and NAdamW for remaining parameters."""
    optimizer_name = "nora_nadamw"
    def __init__(
        self,
        model: nn.Module,
        lr: float = 3e-4,
        weight_decay: float = 0.1,
        nora_lr: float = 5e-3,
        nora_weight_decay: float = 0.0,
        nora_beta: float = 0.95,
        nora_momentum: float = 0.95,
    ):
        groups = partition_optimizer_parameters(model)
        all_params = [
            *groups.nora,
            *groups.nadamw_decay,
            *groups.nadamw_no_decay,
        ]
        if not all_params:
            raise ValueError(
                "Cannot build an optimizer for a model with no trainable parameters"
            )
        super().__init__(all_params, {})
        self.nora = (
            Nora(
                groups.nora,
                lr=nora_lr,
                weight_decay=nora_weight_decay,
                momentum=nora_momentum,
                beta=nora_beta,
            )
            if groups.nora
            else None
        )
        nadamw_groups = []
        if groups.nadamw_decay:
            nadamw_groups.append(
                {"params": groups.nadamw_decay, "weight_decay": weight_decay}
            )
        if groups.nadamw_no_decay:
            nadamw_groups.append(
                {"params": groups.nadamw_no_decay, "weight_decay": 0.0}
            )
        self.nadamw = NAdamW(nadamw_groups, lr=lr) if nadamw_groups else None
        self.param_groups = refresh_param_groups([self.nora, self.nadamw])
    @torch.no_grad()
    def step(self, closure=None):
        return composite_step(
            [opt for opt in (self.nora, self.nadamw) if opt is not None],
            closure,
        )
    def zero_grad(self, set_to_none: bool = True):
        composite_zero_grad(
            [opt for opt in (self.nora, self.nadamw) if opt is not None],
            set_to_none,
        )
    def state_dict(self) -> dict[str, Any]:
        return composite_state_dict({"nora": self.nora, "nadamw": self.nadamw})
    def load_state_dict(self, state_dict: dict[str, Any]):
        if "muon" in state_dict or "adamw" in state_dict:
            raise ValueError(
                "Checkpoint uses muon_adamw state; select optimizer='muon_adamw' "
                "to resume it"
            )
        if "nora" not in state_dict or "nadamw" not in state_dict:
            raise ValueError(
                "Checkpoint optimizer state is not compatible with nora_nadamw"
            )
        saved_nora = state_dict["nora"]
        saved_nadamw = state_dict["nadamw"]
        if (self.nora is None) != (saved_nora is None):
            raise ValueError("Checkpoint Nora parameter groups do not match the model")
        if (self.nadamw is None) != (saved_nadamw is None):
            raise ValueError(
                "Checkpoint NAdamW parameter groups do not match the model"
            )
        if self.nora is not None:
            self.nora.load_state_dict(saved_nora)
        if self.nadamw is not None:
            self.nadamw.load_state_dict(saved_nadamw)
        self.param_groups = refresh_param_groups([self.nora, self.nadamw])
@@ -7,8 +7,9 @@ from astrai.parallel.executor import (
    FSDPExecutor,
    GradientState,
    NoneExecutor,
    broadcast_state_dict,
    create_ref_model,
 )
 from astrai.parallel.module import ColumnParallelLinear, RowParallelLinear
 from astrai.parallel.setup import (
    get_current_device,
    get_rank,
@@ -25,8 +26,6 @@ __all__ = [
    "only_on_rank",
    "setup_parallel",
    "spawn_parallel_fn",
    "RowParallelLinear",
    "ColumnParallelLinear",
    "ExecutorFactory",
    "BaseExecutor",
    "GradientState",
@@ -35,4 +34,6 @@ __all__ = [
    "NoneExecutor",
    "DDPExecutor",
    "FSDPExecutor",
    "create_ref_model",
    "broadcast_state_dict",
 ]
@@ -2,16 +2,21 @@
 import contextlib
 import logging
 import os
 from contextlib import contextmanager
-from typing import Optional, Tuple
+from typing import Any, Callable, Dict, Optional, Tuple
 import torch
 import torch.distributed as dist
 import torch.nn as nn
-from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+from torch.distributed.fsdp import (
    FSDPModule,
    fully_shard,
 )
 from torch.distributed.tensor import DTensor
 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
@@ -19,6 +24,82 @@ from astrai.parallel.setup import get_rank, get_world_size
 logger = logging.getLogger(__name__)
 def broadcast_state_dict(
    state_dict: Optional[Dict[str, torch.Tensor]],
    src: int = 0,
 ) -> Optional[Dict[str, torch.Tensor]]:
    """Broadcast a state_dict from *src* rank to all ranks.
    Tensors stay on their original device (GPU) for the broadcast.
    All ranks must call this collectively.
    On non-distributed runs, returns *state_dict* unchanged.
    """
    if not dist.is_initialized() or dist.get_world_size() == 1:
        return state_dict
    rank = dist.get_rank()
    # Broadcast metadata (keys, shapes, dtypes, device) so non-src ranks
    # can allocate matching empty tensors on the correct device.
    if rank == src:
        device = next(iter(state_dict.values())).device
        metadata = [
            (k, tuple(v.shape), v.dtype, str(device)) for k, v in state_dict.items()
        ]
    else:
        metadata = None
    metadata_list = [metadata]
    dist.broadcast_object_list(metadata_list, src=src)
    metadata = metadata_list[0]
    # Non-src ranks allocate empty tensors with the broadcasted metadata.
    if rank != src:
        state_dict = {
            k: torch.empty(s, dtype=d, device=torch.device(dev))
            for k, s, d, dev in metadata
        }
    # Broadcast each tensor in-place.
    for tensor in state_dict.values():
        dist.broadcast(tensor, src=src)
    return state_dict
 def create_ref_model(
    model_fn: Callable[[], nn.Module],
    executor: Optional["BaseExecutor"] = None,
    model: Optional[nn.Module] = None,
    state_dict: Optional[Dict[str, torch.Tensor]] = None,
    device: Optional[str] = None,
 ) -> Optional[nn.Module]:
    """Create a frozen reference model from executor or state dict.
    In distributed mode (FSDP), ``unwrap_model`` returns ``None`` on
    non-rank-0.  The state_dict is broadcast from rank-0 to all ranks
    so every rank gets a complete copy.
    """
    if state_dict is None and executor is not None and model is not None:
        state_dict = executor.unwrap_model(model)
    # FSDP's unwrap_model returns None on non-rank-0. Broadcast from
    # rank-0 so every rank receives a complete state_dict.
    if executor is not None and executor.use_distributed:
        state_dict = broadcast_state_dict(state_dict)
    if state_dict is None:
        return None
    ref_model = model_fn()
    ref_model.load_state_dict(state_dict)
    ref_model.requires_grad_(False)
    ref_model.eval()
    if device is not None:
        ref_model = ref_model.to(device=device)
    return ref_model
 class GradientState:
    def __init__(self, grad_accum_steps: int = 1):
        self.num_steps = max(grad_accum_steps, 1)
@@ -83,19 +164,28 @@ class BaseExecutor:
    def prepare(
        self,
-        model: nn.Module,
+        model_fn: Callable[[], nn.Module],
-        optimizer: Optional[Optimizer] = None,
+        optimizer_fn: Optional[Callable[[nn.Module], Optimizer]] = None,
-        dataloader: Optional[DataLoader] = None,
+        scheduler_fn: Optional[Callable[[Optimizer], LRScheduler]] = None,
-        scheduler: Optional[LRScheduler] = None,
+        before_wrap: Optional[Callable[[nn.Module], nn.Module]] = None,
-    ) -> Tuple[
+        after_wrap: Optional[Callable[[nn.Module], nn.Module]] = None,
-        nn.Module, Optional[Optimizer], Optional[DataLoader], Optional[LRScheduler]
+    ) -> Tuple[nn.Module, Optional[Optimizer], Optional[LRScheduler]]:
-    ]:
+        model = model_fn()
        if before_wrap is not None:
            model = before_wrap(model)
        model = self._prepare_model(model)
-        if optimizer is not None:
+        if after_wrap is not None:
            model = after_wrap(model)
        optimizer = None
        scheduler = None
        if optimizer_fn is not None:
            optimizer = optimizer_fn(model)
            if scheduler_fn is not None:
                scheduler = scheduler_fn(optimizer)
            optimizer = AccumOptimizer(optimizer, self.gradient_state)
-        if scheduler is not None:
+            if scheduler is not None:
-            scheduler = AccumScheduler(scheduler, self.gradient_state)
+                scheduler = AccumScheduler(scheduler, self.gradient_state)
-        return model, optimizer, dataloader, scheduler
+        return model, optimizer, scheduler
    def _prepare_model(self, model: nn.Module) -> nn.Module:
        return model
@@ -115,8 +205,23 @@ class BaseExecutor:
    def backward(self, loss: torch.Tensor):
        loss.backward()
-    def unwrap_model(self, model: nn.Module) -> nn.Module:
+    def unwrap_model(self, model: nn.Module):
-        return model
+        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:
@@ -130,6 +235,12 @@ class BaseExecutor:
    def grad_accum_steps(self) -> int:
        return self.gradient_state.num_steps
    def clip_grad_norm(self, model: nn.Module, max_norm: float) -> float:
        total_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
        if isinstance(total_norm, torch.Tensor):
            return total_norm.item()
        return total_norm
 class ExecutorFactory(BaseFactory[BaseExecutor]):
    pass
@@ -180,7 +291,7 @@ class DDPExecutor(BaseExecutor):
        if not self.use_distributed:
            logger.warning("DDP backend selected but world_size=1, model not wrapped")
            return model
-        local_rank = get_rank()
+        local_rank = int(os.environ.get("LOCAL_RANK", get_rank()))
        model = DDP(
            model,
            device_ids=[local_rank],
@@ -195,73 +306,123 @@ class DDPExecutor(BaseExecutor):
            return model.no_sync()
        return contextlib.nullcontext()
-    def unwrap_model(self, model: nn.Module) -> nn.Module:
+    def unwrap_model(self, model: nn.Module):
        if isinstance(model, DDP):
-            return model.module
+            return model.module.state_dict()
-        return model
+        return model.state_dict()
@ExecutorFactory.register("fsdp")
 class FSDPExecutor(BaseExecutor):
    """FSDP executor using `torch.distributed.fsdp.fully_shard` (per-module API).
    Wraps each child module individually via ``fully_shard``.
    Skips the root model because ``ABC + Generic[T]`` in the MRO makes
    ``fully_shard``'s dynamic ``__class__`` assignment fail at the CPython level.
    Original ``Parameter`` objects are preserved (as DTensors) — no
    ``FlatParameter``, no ``use_orig_params=True`` hack.
    """
    def __init__(
        self,
        grad_accum_steps: int = 1,
-        process_group=None,
+        mesh: Optional[Any] = None,
-        sharding_strategy=None,
+        mp_policy: Optional[Any] = None,
-        cpu_offload=None,
+        reshard_after_forward: bool = False,
        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,
        use_orig_params: bool = False,
        ignored_states=None,
        device_mesh=None,
    ):
        super().__init__(grad_accum_steps=grad_accum_steps)
-        self._fsdp_kwargs = {
+        self._mesh = mesh
-            k: v
+        self._mp_policy = mp_policy
-            for k, v in dict(
+        self._reshard_after_forward = reshard_after_forward
                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=use_orig_params,
                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())
+        kwargs = dict(
-        model = FSDP(model, device_id=device_id, **self._fsdp_kwargs)
+            mesh=self._mesh,
-        logger.info("Model wrapped with FSDP (world_size=%d)", get_world_size())
+            mp_policy=self._mp_policy,
            reshard_after_forward=self._reshard_after_forward,
        )
        kwargs = {k: v for k, v in kwargs.items() if v is not None}
        for child in model.children():
            if isinstance(child, nn.ModuleList):
                for sub in child:
                    fully_shard(sub, **kwargs)
            else:
                fully_shard(child, **kwargs)
        logger.info(
            "FSDP wrapping applied to %d direct children (root skipped for ABC compat)",
            len(list(model.children())),
        )
        return model
    @contextmanager
    def _no_sync(self, model: nn.Module):
-        if isinstance(model, FSDP):
+        fsdp_modules = [m for m in model.modules() if isinstance(m, FSDPModule)]
-            return model.no_sync()
+        if fsdp_modules:
-        return contextlib.nullcontext()
+            for m in fsdp_modules:
                m.set_requires_gradient_sync(False, recurse=True)
            try:
                yield
            finally:
                for m in fsdp_modules:
                    m.set_requires_gradient_sync(True, recurse=True)
        else:
            yield
-    def unwrap_model(self, model: nn.Module) -> nn.Module:
+    def clip_grad_norm(self, model: nn.Module, max_norm: float) -> float:
-        if self._original_model is not None:
+        if not self.use_distributed:
-            return self._original_model
+            return super().clip_grad_norm(model, max_norm)
-        if isinstance(model, FSDP):
+
-            return model._fsdp_wrapped_module
+        # FSDP params are DTensors (sharded across ranks).
-        return model
+        # torch.nn.utils.clip_grad_norm_ computes LOCAL norm per rank,
        # so we must all-reduce to get the global norm before clipping.
        local_norm = torch.nn.utils.get_total_norm(
            [p.grad for p in model.parameters() if p.grad is not None],
        )
        if isinstance(local_norm, DTensor):
            local_norm = local_norm.to_local()
        total_norm_sq = local_norm**2
        dist.all_reduce(total_norm_sq, op=dist.ReduceOp.SUM)
        total_norm = total_norm_sq.sqrt()
        clip_coef = max_norm / (total_norm + 1e-6)
        clip_coef_clamped = torch.clamp(clip_coef, max=1.0)
        for p in model.parameters():
            if p.grad is not None:
                p.grad.mul_(clip_coef_clamped)
        return total_norm.item()
    def unwrap_model(self, model: nn.Module):
        if not self.use_distributed:
            return model.state_dict()
        # unshard() and full_tensor() are collective ops — all ranks must
        # participate. Non-rank-0 ranks still call them but discard results.
        for module in model.modules():
            if isinstance(module, FSDPModule):
                module.unshard()
        state_dict = model.state_dict()
        result = {}
        for k, v in state_dict.items():
            if isinstance(v, DTensor):
                full = v.full_tensor()
                if get_rank() == 0:
                    result[k] = full
            elif get_rank() == 0:
                result[k] = v
        for module in model.modules():
            if isinstance(module, FSDPModule):
                module.reshard()
        if get_rank() != 0:
            return None
        return result
@@ -1,115 +0,0 @@
 from typing import Dict
 import torch
 import torch.distributed as dist
 import torch.nn as nn
 import torch.nn.functional as F
 from torch import Tensor
 class ParallelModel(nn.Module):
    def __init__(self, process_group: dist.ProcessGroup):
        super().__init__()
        self.process_group = process_group
        self.rank = dist.get_rank(self.process_group)
        self.world_size = dist.get_world_size(self.process_group)
 class RowParallelLinear(ParallelModel):
    def __init__(
        self,
        process_group: dist.ProcessGroup,
        in_features: int,
        out_features: int,
        bias: bool = True,
        reduce_results: bool = True,
    ):
        super().__init__(process_group)
        self.in_features = in_features
        self.out_features = out_features
        self.in_features_per_rank = in_features // self.world_size
        self.reduce_results = reduce_results
        if in_features % self.world_size != 0:
            raise ValueError(
                f"in_features must be divisible by world_size. Got {in_features} and {self.world_size}"
            )
        self.weight = nn.Parameter(torch.empty(out_features, self.in_features_per_rank))
        self.bias = nn.Parameter(torch.zeros(out_features)) if bias else None
    def forward(self, input: Tensor) -> Tensor:
        output = F.linear(input, self.weight)
        if self.reduce_results:
            dist.all_reduce(output, op=dist.ReduceOp.SUM, group=self.process_group)
        if self.bias is not None:
            output += self.bias
        return output
    def load_state_dict(self, state_dict: Dict[str, Tensor]):
        full_weight = state_dict.get("weight")
        full_bias = state_dict.get("bias")
        start_idx = self.rank * self.in_features_per_rank
        end_idx = start_idx + self.in_features_per_rank
        weight_slice = full_weight[:, start_idx:end_idx]
        self.weight.data.copy_(weight_slice)
        if self.bias is not None:
            self.bias.data.copy_(full_bias)
 class ColumnParallelLinear(ParallelModel):
    def __init__(
        self,
        process_group: dist.ProcessGroup,
        in_features: int,
        out_features: int,
        bias: bool = True,
        gather_results: bool = True,
    ):
        super().__init__(process_group)
        self.in_features = in_features
        self.out_features = out_features
        self.out_features_per_rank = out_features // self.world_size
        self.gather_results = gather_results
        if out_features % self.world_size != 0:
            raise ValueError(
                f"out_features must be divisible by world_size. Got {out_features} and {self.world_size}"
            )
        self.weight = nn.Parameter(
            torch.empty(self.out_features_per_rank, self.in_features)
        )
        self.bias = (
            nn.Parameter(torch.zeros(self.out_features_per_rank)) if bias else None
        )
    def forward(self, input: Tensor) -> Tensor:
        output = F.linear(input, self.weight, self.bias)
        if self.gather_results:
            output_list = [torch.empty_like(output) for _ in range(self.world_size)]
            dist.all_gather(output_list, output, group=self.process_group)
            output = torch.cat(output_list, dim=-1)
        return output
    def load_state_dict(self, state_dict: Dict[str, Tensor]):
        full_weight = state_dict.get("weight")
        full_bias = state_dict.get("bias")
        start_idx = self.rank * self.out_features_per_rank
        end_idx = start_idx + self.out_features_per_rank
        weight_slice = full_weight[start_idx:end_idx, :]
        self.weight.data.copy_(weight_slice)
        if self.bias is not None:
            bias_slice = full_bias[start_idx:end_idx]
            self.bias.data.copy_(bias_slice)
@@ -1,12 +1,27 @@
 import logging
 import os
 import signal
 import socket
 import threading
 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
 from astrai.signal_handler import install_early_signal_handlers
 logger = logging.getLogger(__name__)
 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 +45,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 +57,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 +112,7 @@ def only_on_rank(rank, sync=False):
    return decorator
-def wrapper_spawn_func(
+def _run_single_rank(
    rank: int,
    world_size: int,
    backend: str,
@@ -100,20 +122,143 @@ def wrapper_spawn_func(
    func: Callable,
    kwargs: dict,
 ):
-    try:
+    install_early_signal_handlers()
    with setup_parallel(
        rank=rank,
        world_size=world_size,
        local_rank=rank,
        backend=backend,
        master_addr=master_addr,
        master_port=master_port,
        device_type=device_type,
    ):
        func(**kwargs)
 class LaunchStrategy(ABC):
    """Strategy for launching a function in a distributed context."""
    def __init__(
        self,
        world_size: int,
        backend: str,
        master_addr: str,
        master_port: str,
        device_type: str,
        start_method: str,
    ):
        self.world_size = world_size
        self.backend = backend
        self.master_addr = master_addr
        self.master_port = master_port
        self.device_type = device_type
        self.start_method = start_method
    @abstractmethod
    def launch(self, func: Callable, **kwargs):
        raise NotImplementedError
 class TorchrunStrategy(LaunchStrategy):
    """External orchestrator (torchrun, SLURM, K8s) — env vars pre-set."""
    def launch(self, func: Callable, **kwargs):
        install_early_signal_handlers()
        rank = int(os.environ["RANK"])
        world_size = int(os.environ["WORLD_SIZE"])
        local_rank = int(os.environ.get("LOCAL_RANK", rank))
        with setup_parallel(
            rank=rank,
            world_size=world_size,
-            backend=backend,
+            local_rank=local_rank,
-            master_addr=master_addr,
+            backend=self.backend,
-            master_port=master_port,
+            master_addr=os.environ.get("MASTER_ADDR", self.master_addr),
-            device_type=device_type,
+            master_port=os.environ.get("MASTER_PORT", self.master_port),
            device_type=self.device_type,
        ):
            func(**kwargs)
-    except Exception as e:
+
-        print(f"Error in rank {rank}: {e}")
+class LocalStrategy(LaunchStrategy):
-        raise
+    """Local launcher — single-process or mp.start_processes."""
    def launch(self, func: Callable, **kwargs):
        args = (
            self.world_size,
            self.backend,
            self.master_addr,
            self.master_port,
            self.device_type,
            func,
            kwargs,
        )
        if self.world_size == 1:
            _run_single_rank(0, *args)
            return
        install_early_signal_handlers()
        ctx = mp.start_processes(
            _run_single_rank,
            args=args,
            nprocs=self.world_size,
            start_method=self.start_method,
            join=False,
        )
        parent_stop = threading.Event()
        original_handlers = {}
        def _parent_handler(signum, frame):
            sig = signal.Signals(signum)
            logger.warning(
                "Parent (pid=%d) received %s, forwarding to children...",
                os.getpid(),
                sig.name,
            )
            parent_stop.set()
            for p in ctx.processes:
                if p.is_alive():
                    p.terminate()
        for sig in (signal.SIGTERM, signal.SIGINT):
            prev = signal.signal(sig, _parent_handler)
            if prev not in (signal.SIG_DFL, signal.SIG_IGN, None, _parent_handler):
                original_handlers[sig] = prev
        try:
            while not ctx.join() and not parent_stop.is_set():
                pass
        except BaseException:
            logger.warning(
                "Parent received unexpected exception, terminating children..."
            )
            for p in ctx.processes:
                if p.is_alive():
                    p.terminate()
            raise
        finally:
            for sig, handler in original_handlers.items():
                signal.signal(sig, handler)
            for p in ctx.processes:
                p.join()
            ctx.join()
 def _detect_launcher() -> str:
    """Detect the distributed launcher from environment.
    Returns one of: "torchelastic", "torchrun", "external", "local".
    """
    if dist.is_torchelastic_launched():
        return "torchelastic"
    if "LOCAL_WORLD_SIZE" in os.environ:
        return "torchrun"
    if "RANK" in os.environ and "WORLD_SIZE" in os.environ:
        return "external"
    return "local"
 def spawn_parallel_fn(
@@ -121,46 +266,20 @@ def spawn_parallel_fn(
    world_size: int,
    backend: str = "nccl",
    master_addr: str = "localhost",
-    master_port: str = "29500",
+    master_port: Optional[str] = None,
    device_type: str = "cuda",
    start_method: str = "spawn",
    **kwargs,
 ):
-    # clear environment variables
+    if master_port is None:
-    for key in [
+        master_port = find_free_port()
-        "MASTER_ADDR",
+    launcher = _detect_launcher()
-        "MASTER_PORT",
+    if launcher in ("torchelastic", "torchrun", "external"):
-        "RANK",
+        strategy = TorchrunStrategy(
-        "WORLD_SIZE",
+            world_size, backend, master_addr, master_port, device_type, start_method
-        "LOCAL_RANK",
+        )
-        "LOCAL_DEVICE",
+    else:
-    ]:
+        strategy = LocalStrategy(
-        if key in os.environ:
+            world_size, backend, master_addr, master_port, device_type, start_method
-            del os.environ[key]
+        )
-
+    strategy.launch(func, **kwargs)
    if world_size == 1:
        device_id = torch.device(device_type, 0)
        os.environ["LOCAL_RANK"] = "0"
        os.environ["WORLD_SIZE"] = "1"
        os.environ["LOCAL_DEVICE"] = str(device_id)
        func(**kwargs)
        return
    wrapper_spawn_func_args = (
        world_size,
        backend,
        master_addr,
        master_port,
        device_type,
        func,
        kwargs,
    )
    mp.start_processes(
        wrapper_spawn_func,
        args=wrapper_spawn_func_args,
        nprocs=world_size,
        start_method=start_method,
        join=True,
    )
@@ -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,542 @@
 """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_sections_batch(
        self,
        items: list[dict],
        sections: list,
        config,
        tokenizer,
        *,
        is_top_level=False,
        filter_text=True,
    ):
        """Render and tokenize a group of records with batched Rust tokenization."""
        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
        )
        plans: list[list[tuple[str, str, bool]]] = []
        for item in items:
            plan: list[tuple[str, str, bool]] = []
            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:
                    messages = item.get(field)
                    if not isinstance(messages, list) or not messages:
                        continue
                    for msg in messages:
                        role = msg.get("role", "")
                        rendered = tokenizer.apply_chat_template(
                            [msg], tokenize=False, add_generation_prompt=False
                        )
                        plan.append(
                            (rendered, _resolve_action(action, role, config), False)
                        )
                else:
                    text = str(item.get(field, ""))
                    if not text.strip():
                        continue
                    if is_text_config and filter_text:
                        pp = config.preprocessing
                        if pp.min_chars > 0 and len(text) < pp.min_chars:
                            continue
                        if len(text) > pp.max_chars:
                            continue
                    plan.append((text, action, add_special))
                first_section = False
            plans.append(plan)
        encoded: dict[tuple[int, int], list[int]] = {}
        for add_special in (False, True):
            refs = [
                (item_idx, unit_idx, text)
                for item_idx, plan in enumerate(plans)
                for unit_idx, (text, _, add) in enumerate(plan)
                if add == add_special
            ]
            if not refs:
                continue
            ids_batch = tokenizer.encode(
                [text for _, _, text in refs], add_special_tokens=add_special
            )
            for (item_idx, unit_idx, _), ids in zip(refs, ids_batch):
                encoded[(item_idx, unit_idx)] = ids
        outputs = []
        max_len = config.preprocessing.max_seq_len
        for item_idx, plan in enumerate(plans):
            all_ids = []
            loss_mask = []
            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)
            for unit_idx, (_, action, _) in enumerate(plan):
                ids = encoded[(item_idx, unit_idx)]
                all_ids.extend(ids)
                loss_mask.extend([1 if action == "train" else 0] * len(ids))
            all_ids = all_ids[:max_len]
            loss_mask = loss_mask[: len(all_ids)]
            if not all_ids or (is_top_level and has_template and len(all_ids) <= 1):
                outputs.append((None, None))
            else:
                outputs.append((all_ids, loss_mask))
        return outputs
    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
    def process_list_field_batch(self, items, sections, config, tokenizer):
        per_item_ids = [[] for _ in items]
        per_item_masks = [[] for _ in items]
        for sec in sections:
            wrappers = []
            owners = []
            field = sec["field"]
            for item_idx, item in enumerate(items):
                values = item.get(field)
                if not isinstance(values, list):
                    continue
                for val in values:
                    if sec.get("template", False) and not isinstance(val, list):
                        continue
                    wrappers.append({field: val if isinstance(val, list) else str(val)})
                    owners.append(item_idx)
            rendered = self.process_sections_batch(
                wrappers,
                [sec],
                config,
                tokenizer,
                is_top_level=False,
                filter_text=False,
            )
            for owner, (ids, mask) in zip(owners, rendered):
                if ids:
                    per_item_ids[owner].append(ids)
                    per_item_masks[owner].append(mask)
        return [
            (ids, masks) if ids else (None, None)
            for ids, masks in zip(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]: ...
    def build_batch(self, items: list[dict], config, tokenizer) -> list[Optional[dict]]:
        return [self.build(item, config, tokenizer) for item in items]
 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
    def build_batch(self, items, config, tokenizer):
        sections = config.input.sections
        if not sections:
            return [None] * len(items)
        rendered = self.renderer.process_sections_batch(
            items, sections, config, tokenizer, is_top_level=True
        )
        results = []
        for item, (ids, mask) in zip(items, rendered):
            if ids is None:
                results.append(None)
                continue
            result = {
                "sequence": ids,
                "domain": _extract_domain(item, config.output.domain_key),
            }
            if not all(m == 1 for m in mask):
                result["loss_mask"] = mask
            results.append(result)
        return results
@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 = {}
        required_outputs = {
            output_key
            for output_key, spec in sources_spec.items()
            if spec.get("sections")
        }
        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
                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
                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
        if not required_outputs or not required_outputs.issubset(result):
            return None
        result["domain"] = _extract_domain(item, config.output.domain_key)
        return result
    def build_batch(self, items, config, tokenizer):
        sources_spec = getattr(config.input, "sources", None)
        if not sources_spec:
            return [None] * len(items)
        results = [{} for _ in items]
        required_outputs = {
            output_key
            for output_key, spec in sources_spec.items()
            if spec.get("sections")
        }
        for output_key, spec in sources_spec.items():
            sections = spec.get("sections", [])
            if not sections:
                continue
            if self.renderer.is_value_section(sections):
                for item, result in zip(items, results):
                    value = self.renderer.extract_raw_value(item, sections)
                    if value is not None:
                        result[output_key] = value
                continue
            mask_key = spec.get("mask_key", f"{output_key}_mask")
            if spec.get("list_field", False):
                rendered = self.renderer.process_list_field_batch(
                    items, sections, config, tokenizer
                )
            else:
                rendered = self.renderer.process_sections_batch(
                    items, sections, config, tokenizer, is_top_level=True
                )
            for result, (ids, mask) in zip(results, rendered):
                if ids is None:
                    continue
                result[output_key] = ids
                if spec.get("list_field", False) or not all(m == 1 for m in mask):
                    result[mask_key] = mask
                elif "mask_key" in spec:
                    result[mask_key] = mask
        return [
            ({**result, "domain": _extract_domain(item, config.output.domain_key)})
            if required_outputs and required_outputs.issubset(result)
            else None
            for item, result in zip(items, results)
        ]
@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)
    def build_batch(self, items, config, tokenizer):
        sources_spec = getattr(config.input, "sources", None)
        if sources_spec:
            return self._multi.build_batch(items, config, tokenizer)
        return self._single.build_batch(items, 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,277 @@
 """Config-driven JSONL preprocessing pipeline.
 Composes a :class:`BaseMaskBuilder` (selected by ``input.type``) with
 sharding and flush to ``.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,
    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 transform_batch(self, items: list[dict]) -> list[Optional[dict]]:
        return self.mask_builder.build_batch(items, 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
        progress = tqdm.tqdm(desc="Tokenizing", unit="docs", mininterval=0.5)
        stop = False
        for items in self._iter_batches(pp.batch_size):
            progress.update(len(items))
            try:
                results = self.transform_batch(items)
            except Exception:
                logger.warning(
                    "Failed to process batch, retrying records individually",
                    exc_info=True,
                )
                results = []
                for item in items:
                    try:
                        results.append(self.transform(item))
                    except Exception:
                        logger.warning(
                            "Failed to process item, skipping", exc_info=True
                        )
                        results.append(None)
            for result in results:
                if pp.max_items and count >= pp.max_items:
                    stop = True
                    break
                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 stop:
                break
        progress.close()
        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 _iter_batches(self, batch_size: int):
        batch_size = max(1, batch_size)
        batch = []
        for item in self._iter_items():
            batch.append(item)
            if len(batch) >= batch_size:
                yield batch
                batch = []
        if batch:
            yield batch
    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,56 @@
 """Storage writer strategies for pipeline output.
 The :class:`StoreWriter` abstraction decouples the pipeline from the
 concrete storage format (bin).  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
 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
@@ -0,0 +1,41 @@
 """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,
    save_bin,
 )
 __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",
    "save_bin",
 ]
@@ -1,9 +1,11 @@
 """Model checkpoint serialization helpers."""
 import io
 import json
 import time
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Any, Dict, Union
+from typing import Any, Dict, Optional, Union
 import safetensors.torch as st
 import torch
@@ -136,7 +138,7 @@ def load_state_dict(path: Union[str, Path], broadcast: bool = False) -> dict:
 class Checkpoint:
    state_dict: Dict[str, Any] = field(default_factory=dict)
    epoch: int = 0
-    iteration: 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)
@@ -145,12 +147,9 @@ class Checkpoint:
        save_path = Path(save_dir)
        save_path.mkdir(parents=True, exist_ok=True)
        if get_rank() != 0:
            return
        meta = {
            "epoch": self.epoch,
-            "iteration": self.iteration,
+            "consumed_samples": self.consumed_samples,
            "timestamp": time.strftime("%Y-%m-%dT%H:%M:%S"),
            **self.meta,
        }
@@ -176,7 +175,27 @@ class Checkpoint:
        return cls(
            state_dict=state_dict,
            epoch=meta.get("epoch", 0),
-            iteration=meta.get("iteration", 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,82 @@
 """Dataset storage serialization helpers (memory-mapped binary)."""
 import json
 import os
 from typing import Any, Dict, List, Optional
 import numpy as np
 import torch
 from torch import Tensor
 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 JSONL 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 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="c",
            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 (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
@@ -0,0 +1,53 @@
 import logging
 import os
 import signal
 import threading
 logger = logging.getLogger(__name__)
 _early_stop = threading.Event()
 _active_context = None
 def _early_handler(signum: int, frame):
    sig = signal.Signals(signum)
    logger.warning(
        "Received %s (pid=%d), requesting graceful training stop...",
        sig.name,
        os.getpid(),
    )
    _early_stop.set()
    if _active_context is not None:
        _active_context.request_stop()
 def install_early_signal_handlers():
    for sig in (signal.SIGTERM, signal.SIGINT):
        signal.signal(sig, _early_handler)
    _unblock_signals()
 def _unblock_signals():
    try:
        mask = signal.pthread_sigmask(signal.SIG_BLOCK, set())
        blocked = {signal.SIGTERM, signal.SIGINT} & mask
        if blocked:
            signal.pthread_sigmask(signal.SIG_UNBLOCK, blocked)
    except (AttributeError, OSError):
        pass
 def register_signal_handlers(context):
    global _active_context
    _active_context = context
    for sig in (signal.SIGTERM, signal.SIGINT):
        signal.signal(sig, _early_handler)
    if _early_stop.is_set():
        context.request_stop()
        logger.warning("Signal was received during initialization, stopping...")
 def unregister_signal_handlers():
    global _active_context
    _active_context = None
    _early_stop.clear()
@@ -1,8 +1,10 @@
 from astrai.tokenize.chat_template import ChatTemplate, MessageType
-from astrai.tokenize.tokenizer import AutoTokenizer
+from astrai.tokenize.tokenizer import AutoTokenizer, Message, Messages
 __all__ = [
    "AutoTokenizer",
    "ChatTemplate",
    "MessageType",
    "Message",
    "Messages",
 ]
@@ -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,51 @@ 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.  :meth:`__getstate__` drops the cached template so
        that pickle serialises only ``template_str``; each worker rebuilds
        the cache on first render.
        """
        return Template(self.template_str)
    def __getstate__(self) -> Dict[str, Any]:
        """Exclude the cached Jinja2 template from pickling.
        ``Template.root_render_func`` is a dynamically generated closure
        that cannot be pickled by reference.  Dropping ``_compiled`` here
        lets :class:`cached_property` rebuild it on first access after
        unpickle.
        """
        state = self.__dict__.copy()
        state.pop("_compiled", None)
        return state
    def __setstate__(self, state: Dict[str, Any]) -> None:
        self.__dict__.update(state)
    @classmethod
    def from_string(
@@ -43,7 +69,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 +99,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)
@@ -10,12 +10,16 @@ from tokenizers import Tokenizer
 from astrai.tokenize.chat_template import ChatTemplate
 Message = Dict[str, str]
 """Single chat message with ``role`` and ``content`` keys."""
 Messages = List[Message]
 """Single conversation — a list of messages."""
 class AutoTokenizer:
    """Base tokenizer class with automatic loading support"""
    TOKENIZER_CLASSES = {}  # Registry for auto-loading
    def __init__(
        self,
        path: Optional[Union[str, Path]] = None,
@@ -102,17 +106,6 @@ class AutoTokenizer:
        with open(save_path / "tokenizer_config.json", "w", encoding="utf-8") as f:
            json.dump(config, f, ensure_ascii=False, indent=2)
    @classmethod
    def register_tokenizer(cls, name: str, tokenizer_class: type):
        """
        Register a new tokenizer class.
        Args:
            name: Name to register the tokenizer class under
            tokenizer_class: The tokenizer class to register
        """
        cls.TOKENIZER_CLASSES[name] = tokenizer_class
    def encode(
        self,
        tokens: Union[str, List[str]],
@@ -120,7 +113,16 @@ class AutoTokenizer:
        is_pretokenized: bool = False,
        add_special_tokens: bool = True,
    ) -> List:
-        """Encode text to tokens or token IDs."""
+        """Encode text to token IDs.
        Accepts both single strings and batches:
        - ``encode("hello")`` → ``[123, 456]``
        - ``encode(["hello", "world"])`` → ``[[123, 456], [789]]``
        Batches are tokenised in parallel via the Rust backend's
        ``encode_batch`` (uses all available CPU cores).
        """
        if self._tokenizer is None:
            raise RuntimeError(
                "Tokenizer not initialized. Load or create a tokenizer first."
@@ -133,15 +135,13 @@ class AutoTokenizer:
                add_special_tokens=add_special_tokens,
            )
            return encoded.ids if out_ids else encoded.tokens
-        else:
+
-            encoded_list = self._tokenizer.encode_batch(
+        encoded_list = self._tokenizer.encode_batch(
-                tokens,
+            tokens,
-                is_pretokenized=is_pretokenized,
+            is_pretokenized=is_pretokenized,
-                add_special_tokens=add_special_tokens,
+            add_special_tokens=add_special_tokens,
-            )
+        )
-            return [
+        return [encoded.ids if out_ids else encoded.tokens for encoded in encoded_list]
                encoded.ids if out_ids else encoded.tokens for encoded in encoded_list
            ]
    def decode(self, tokens: List[int], skip_special_tokens: bool = True) -> str:
        """Decode token IDs to text."""
@@ -164,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 = []
@@ -220,45 +227,63 @@ class AutoTokenizer:
    def apply_chat_template(
        self,
-        messages: List[Dict[str, str]],
+        messages: Union[Messages, List[Messages]],
        system_prompt: Optional[str] = None,
        tokenize: bool = True,
        add_generation_prompt: bool = True,
        **kwargs,
-    ) -> Union[str, List[int]]:
+    ) -> Union[str, List[int], List[str], List[List[int]]]:
-        """
+        """Apply the chat template and optionally tokenize.
-        Apply the chat template to messages and optionally tokenize the result.
+
        Accepts both single conversations and batches:
        - ``apply_chat_template([msg1, msg2])`` → ``"..."`` or ``[ids]``
        - ``apply_chat_template([[msg1, msg2], [msg3]])`` → ``["..", ".."]``
          or ``[[ids], [ids]]``
        Batches render each conversation list and tokenise all at once via
        :meth:`encode` (``List[str]`` → Rust parallel ``encode_batch``).
        Args:
-            messages: List of message dicts with 'role' and 'content'.
+            messages: Single conversation (``Messages``) or batch of
-            system_prompt: Optional system prompt string (auto-converted to first message).
+                conversations (``BatchMessages``).
            system_prompt: Optional system prompt prepended (single mode only).
            tokenize: Whether to return token IDs (True) or raw string (False).
-            add_generation_prompt: Whether to add the generation prompt (default: True).
+            add_generation_prompt: Whether to add the generation prompt.
-            **kwargs: Additional variables to pass to the template.
+            **kwargs: Additional template variables.
        Returns:
-            Either the rendered string or list of token IDs.
+            Single mode: ``str`` or ``List[int]``.
-
+            Batch mode: ``List[str]`` or ``List[List[int]]``.
        Raises:
            RuntimeError: If chat template is not set.
        """
        if self._chat_template is None:
            raise RuntimeError(
                "Chat template not set. Use set_chat_template() to set a template first."
            )
-        # Auto-convert system_prompt to first message if provided
+        is_batch = bool(messages) and isinstance(messages[0], list)
        if is_batch:
            rendered = [
                self._chat_template.render(
                    messages=msgs,
                    add_generation_prompt=add_generation_prompt,
                    **kwargs,
                )
                for msgs in messages
            ]
            if tokenize:
                return self.encode(rendered)  # List[str] → batch encode
            return rendered
        # Single conversation
        if system_prompt:
            messages = [{"role": "system", "content": system_prompt}] + list(messages)
        # Render the template
        rendered = self._chat_template.render(
            messages=messages,
            add_generation_prompt=add_generation_prompt,
            **kwargs,
        )
        if tokenize:
            return self.encode(rendered)
        return rendered
@@ -1,4 +1,3 @@
 from astrai.trainer.optim import Muon
 from astrai.trainer.schedule import BaseScheduler, SchedulerFactory
 from astrai.trainer.strategy import BaseStrategy, StrategyFactory
 from astrai.trainer.train_callback import (
@@ -10,8 +9,6 @@ from astrai.trainer.trainer import Trainer
 __all__ = [
    # Main trainer
    "Trainer",
    # Optimizer
    "Muon",
    # Strategy factory
    "StrategyFactory",
    "BaseStrategy",
@@ -1,42 +1,70 @@
-from typing import Any, Callable, Dict
+from typing import Dict
 import torch
 import torch.nn as nn
-def _grad_stat(
+def grad_norm(model: nn.Module, per_param: bool = False) -> float | Dict[str, float]:
-    model: nn.Module, fn: Callable[[torch.Tensor], Any], default: Any
+    grads = [p.grad.detach() for p in model.parameters() if p.grad is not None]
-) -> dict:
+    if not grads:
-    results = {}
+        return 0.0
-    for name, param in model.named_parameters():
+
-        results[name] = default
+    total_sq = torch.stack([g.pow(2).sum() for g in grads]).sum()
-        if param.grad is not None:
+    if per_param:
-            results[name] = fn(param.grad.data)
+        norms = {}
-    return results
+        for name, param in model.named_parameters():
            if param.grad is not None:
                norms[name] = param.grad.norm(2).item()
            else:
                norms[name] = 0.0
        norms["total"] = total_sq.sqrt().item()
        return norms
    return total_sq.sqrt().item()
-def grad_norm(model: nn.Module, norm_type: int = 2) -> Dict[str, float]:
+class GradSNRTracker:
-    return _grad_stat(model, lambda g: g.norm(norm_type).item(), 0.0)
+    """Track gradient signal-to-noise ratio via EMA of first/second moments.
    SNR = E[g]^2 / Var(g) = E[g]^2 / (E[g^2] - E[g]^2)
-def grad_std(model: nn.Module) -> Dict[str, float]:
+    The tracker accumulates per-parameter EMA moments across optimizer steps.
-    return _grad_stat(model, lambda g: g.std().item(), 0.0)
+    Call ``update`` after backward (before ``optimizer.step``) and read
    ``snr`` to get the aggregate SNR across all parameters.
    """
    def __init__(self, beta: float = 0.999, eps: float = 1e-8):
        self.beta = beta
        self.eps = eps
        self._first: Dict[int, torch.Tensor] = {}
        self._second: Dict[int, torch.Tensor] = {}
-def grad_max(model: nn.Module) -> Dict[str, float]:
+    @torch.no_grad()
-    return _grad_stat(model, lambda g: g.max().item(), -float("inf"))
+    def update(self, model: nn.Module) -> None:
        beta = self.beta
        for param in model.parameters():
            if param.grad is None:
                continue
            pid = id(param)
            g = param.grad.detach()
            if pid not in self._first:
                self._first[pid] = g.clone()
                self._second[pid] = g.pow(2).clone()
            else:
                self._first[pid].mul_(beta).add_(g, alpha=1 - beta)
                self._second[pid].mul_(beta).addcmul_(g, g, value=1 - beta)
-
+    @property
-def grad_min(model: nn.Module) -> Dict[str, float]:
+    def snr(self) -> float:
-    return _grad_stat(model, lambda g: g.min().item(), float("inf"))
+        if not self._first:
-
+            return 0.0
-
+        total_signal = 0.0
-def grad_mean(model: nn.Module) -> Dict[str, float]:
+        total_noise = 0.0
-    return _grad_stat(model, lambda g: g.mean().item(), 0.0)
+        for m, v in zip(self._first.values(), self._second.values()):
-
+            signal = m.pow(2).sum().item()
-
+            noise = (v - m.pow(2)).clamp(min=0).sum().item()
-def grad_nan_num(model: nn.Module) -> Dict[str, int]:
+            total_signal += signal
-    return _grad_stat(model, lambda g: g.isnan().sum().item(), 0)
+            total_noise += noise
        return total_signal / (total_noise + self.eps)
 def ctx_get_loss(ctx):
@@ -52,24 +80,31 @@ def ctx_get_val_loss(ctx):
 def ctx_get_grad_norm(ctx):
-    return grad_norm(ctx.model)
+    return ctx.grad_norm
-def ctx_get_grad_std(ctx):
+def ctx_get_grad_snr(ctx):
-    return grad_std(ctx.model)
+    tracker = getattr(ctx, "grad_snr_tracker", None)
    if tracker is None:
        return None
    return tracker.snr
-def ctx_get_grad_max(ctx):
+def ctx_get_moe_aux_loss(ctx):
-    return grad_max(ctx.model)
+    return ctx.strategy._moe_metrics.get("aux_loss")
-def ctx_get_grad_min(ctx):
+def ctx_get_router_entropy(ctx):
-    return grad_min(ctx.model)
+    return ctx.strategy._moe_metrics.get("router_entropy")
-def ctx_get_grad_mean(ctx):
+def ctx_get_dead_expert_fraction(ctx):
-    return grad_mean(ctx.model)
+    return ctx.strategy._moe_metrics.get("dead_expert_fraction")
-def ctx_get_grad_nan_num(ctx):
+def ctx_get_load_imbalance_mean(ctx):
-    return grad_nan_num(ctx.model)
+    return ctx.strategy._moe_metrics.get("load_imbalance_mean")
 def ctx_get_load_imbalance_max(ctx):
    return ctx.strategy._moe_metrics.get("load_imbalance_max")
@@ -1,143 +0,0 @@
 import torch
 from torch.optim import Optimizer
 def _zeropower_via_newtonschulz(G: torch.Tensor, steps: int = 5):
    assert G.ndim == 2
    X = G
    scale = max(1, G.size(0) / G.size(1)) ** 0.5
    X = X / (X.norm() + 1e-7) * scale
    if steps == 0:
        return X
    a, b, c = (3.4445, -4.7750, 2.0315)
    for _ in range(steps):
        A = X @ X.T
        B = A @ X
        X = a * X + b * B + c * (A @ B)
    return X
 class Muon(Optimizer):
    def __init__(
        self,
        params,
        lr: float = 2e-3,
        momentum: float = 0.95,
        weight_decay: float = 0.0,
        nesterov: bool = True,
        ns_steps: int = 5,
        adamw_lr: float = None,
        adamw_betas: tuple = (0.9, 0.95),
        adamw_eps: float = 1e-8,
        adamw_wd: float = 0.0,
    ):
        defaults = dict(
            lr=lr,
            momentum=momentum,
            weight_decay=weight_decay,
            nesterov=nesterov,
            ns_steps=ns_steps,
            adamw_lr=adamw_lr if adamw_lr is not None else lr * 0.1,
            adamw_betas=adamw_betas,
            adamw_eps=adamw_eps,
            adamw_wd=adamw_wd,
        )
        super().__init__(params, defaults)
    @torch.no_grad()
    def step(self, closure=None):
        loss = None
        if closure is not None:
            with torch.enable_grad():
                loss = closure()
        for group in self.param_groups:
            params_2d, params_1d = [], []
            grads_2d, grads_1d = [], []
            for p in group["params"]:
                if p.grad is None:
                    continue
                if p.grad.is_sparse:
                    raise RuntimeError("Muon does not support sparse gradients")
                if p.ndim >= 2:
                    params_2d.append(p)
                    grads_2d.append(p.grad)
                else:
                    params_1d.append(p)
                    grads_1d.append(p.grad)
            if params_2d:
                self._muon_update_foreach(params_2d, grads_2d, group)
            if params_1d:
                self._adamw_update_foreach(params_1d, grads_1d, group)
        return loss
    def _muon_update_foreach(self, params_2d, grads_2d, group):
        lr = group["lr"]
        momentum = group["momentum"]
        wd = group["weight_decay"]
        nesterov = group["nesterov"]
        ns_steps = group["ns_steps"]
        if wd != 0:
            torch._foreach_mul_(params_2d, 1 - lr * wd)
        if nesterov:
            grads_2d = torch._foreach_add(grads_2d, params_2d, alpha=wd)
        bufs = []
        for p, grad in zip(params_2d, grads_2d):
            state = self.state[p]
            if "momentum_buffer" not in state:
                state["momentum_buffer"] = torch.zeros_like(grad)
            bufs.append(state["momentum_buffer"])
        torch._foreach_lerp_(bufs, grads_2d, 1 - momentum)
        for p, buf in zip(params_2d, bufs):
            update = _zeropower_via_newtonschulz(buf, steps=ns_steps)
            scale = max(1, p.size(0) / p.size(1)) ** 0.5
            p.add_(update, alpha=-lr * scale)
    def _adamw_update_foreach(self, params_1d, grads_1d, group):
        lr = group["adamw_lr"]
        betas = group["adamw_betas"]
        eps = group["adamw_eps"]
        wd = group["adamw_wd"]
        steps: list[int] = []
        exp_avgs, exp_avg_sqs = [], []
        has_state = []
        for p in params_1d:
            state = self.state[p]
            if not state:
                state["step"] = 0
                state["exp_avg"] = torch.zeros_like(p)
                state["exp_avg_sq"] = torch.zeros_like(p)
                has_state.append(False)
            else:
                has_state.append(True)
            state["step"] += 1
            steps.append(state["step"])
            exp_avgs.append(state["exp_avg"])
            exp_avg_sqs.append(state["exp_avg_sq"])
        beta1, beta2 = betas
        torch._foreach_lerp_(exp_avgs, grads_1d, 1 - beta1)
        grads_sq = torch._foreach_mul(grads_1d, grads_1d)
        torch._foreach_lerp_(exp_avg_sqs, grads_sq, 1 - beta2)
        bias_correction1 = [1 - beta1**s for s in steps]
        bias_correction2 = [1 - beta2**s for s in steps]
        if wd != 0:
            torch._foreach_mul_(params_1d, 1 - lr * wd)
        exp_avg_corrected = torch._foreach_div(exp_avgs, bias_correction1)
        denom = torch._foreach_div(exp_avg_sqs, bias_correction2)
        denom = torch._foreach_sqrt(denom)
        torch._foreach_add_(denom, eps)
        torch._foreach_addcdiv_(params_1d, exp_avg_corrected, denom, value=-lr)
@@ -0,0 +1,421 @@
 """Online rollout runner for RL training.
 Provides:
 - :class:`RawRollout` — generation output container (no reward yet)
 - :class:`RolloutResult` — a :class:`RawRollout` with rewards attached
 - :class:`BaseRewardModel` — pluggable reward interface
 - :class:`RolloutGenerator` — KV-cache-backed generation of grouped
  responses + decoding (no reward); delegates the generation loop to
  :class:`~astrai.inference.core.scheduler.InferenceScheduler.run_batch`
  so rollout and the production inference server share one code path
 - :class:`RolloutRunner` — orchestrates generation + scoring with a
  step-driven cache; its ``__call__`` returns ``(RolloutResult, is_fresh)``
  so callers do not need to rely on object identity to detect refreshes.
 """
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, field
 from typing import Dict, List, Optional, Tuple
 import torch
 from torch import Tensor
 from astrai.inference.core.scheduler import InferenceScheduler
@dataclass(kw_only=True)
 class RawRollout:
    """Generation output before reward scoring.
    Produced by :class:`RolloutGenerator`; consumed by :class:`RolloutRunner`
    to assemble a :class:`RolloutResult` once rewards are attached.
    Fields are designed to cover all common RL algorithms:
    GRPO, PPO, Online DPO, Rejection Sampling, etc.
    Fields:
        prompts: Tokenized prompts, shape ``[B, P_len]``.
        prompt_mask: Boolean mask for real prompt tokens, shape ``[B, P_len]``.
        responses: Generated response token IDs, shape ``[B, G, R_max]``.
        response_mask: Boolean mask for real (non-pad) response tokens,
            shape ``[B, G, R_max]``.
        logprobs_old: Per-token log-probs under the behaviour policy,
            shape ``[B, G, R_max]``.
        prompt_texts: Decoded prompt strings (for reward models that
            need text).
        response_texts: Decoded response strings, shape ``[B, G]``
            (for reward models).
    """
    prompts: Tensor
    prompt_mask: Tensor
    responses: Tensor
    response_mask: Tensor
    logprobs_old: Tensor
    prompt_texts: List[str] = field(default_factory=list)
    response_texts: List[List[str]] = field(default_factory=list)
@dataclass(kw_only=True)
 class RolloutResult(RawRollout):
    """A :class:`RawRollout` with reward scoring attached.
    Produced by :class:`RolloutRunner` once the :class:`BaseRewardModel`
    has scored the decoded responses.
    Fields:
        rewards: Reward per response, shape ``[B, G]``.
    """
    rewards: Tensor
 class BaseRewardModel(ABC):
    """Pluggable reward model interface.
    Subclasses should implement ``score()`` to return a ``[B, G]`` float
    tensor of rewards.  Implementations can be:
    * A loaded reward model (e.g. ArmoRM, Skywork-Reward)
    * An external API call
    * A rule-based function (format, length, keyword matching)
    """
    @abstractmethod
    def score(self, prompts: List[str], responses: List[List[str]]) -> Tensor:
        """Score each generated response.
        Args:
            prompts: Raw prompt strings, length ``B``.
            responses: Generated response strings, shape ``[B, G]``.
        Returns:
            Float tensor of shape ``[B, G]``.
        """
        ...
 _PAD = 0
 class RolloutGenerator:
    """Pure generation + decoding for a group of responses per prompt.
    Delegates the prefill/decode loop to
    :meth:`~astrai.inference.core.scheduler.InferenceScheduler.run_batch`,
    which uses a real KV cache (no O(n²) recompute).  Has no dependency
    on any reward model; can be reused in isolation for offline
    generation, qualitative sampling, or eval pipelines.
    """
    def __init__(
        self,
        scheduler: InferenceScheduler,
        tokenizer,
        max_tokens: int = 1024,
        group_size: int = 8,
        temperature: float = 1.0,
        top_k: int = 0,
        top_p: float = 1.0,
        frequency_penalty: float = 0.0,
        rep_window: int = 64,
    ):
        self.scheduler = scheduler
        self.tokenizer = tokenizer
        self.max_tokens = max_tokens
        self.group_size = group_size
        self.temperature = temperature
        self.top_k = top_k
        self.top_p = top_p
        self.frequency_penalty = frequency_penalty
        self.rep_window = rep_window
    @torch.no_grad()
    def generate(self, batch: Dict) -> RawRollout:
        """Expand prompts by ``group_size`` and generate one response each.
        Accepted batch formats (per sample, repeated B times):
        - **messages**: ``{"messages": [{"role": "user", "content": "..."}, ...]}``
        - **instruction + input + output**: ``{"instruction": "...",
          "input": "...", "output": "..."}`` — mapped to ``system`` /
          ``user`` / ``assistant`` messages; ``input`` and ``output``
          are optional and skipped when empty.
        Both are rendered through the tokenizer's chat template with
        ``add_generation_prompt=True`` so rollout prompts match the
        format the policy was SFT-trained on.
        """
        model = self.scheduler._executor.model
        was_training = model.training
        model.eval()
        try:
            return self._generate_eval(batch)
        finally:
            model.train(was_training)
    def _generate_eval(self, batch: Dict) -> RawRollout:
        prompt_texts, flat_prompt_ids = self._prepare_prompts(batch)
        B = len(prompt_texts)
        G = self.group_size
        # Re-expand flat list to G copies per prompt for run_batch.
        expanded_prompt_ids: List[List[int]] = []
        for ids in flat_prompt_ids:
            expanded_prompt_ids.extend([list(ids)] * G)
        results = self.scheduler.run_batch(
            expanded_prompt_ids,
            max_tokens=self.max_tokens,
            temperature=self.temperature,
            top_k=self.top_k,
            top_p=self.top_p,
            frequency_penalty=self.frequency_penalty,
            rep_window=self.rep_window,
            return_logprobs=True,
        )
        if len(results) != B * G:
            raise RuntimeError(
                f"Rollout scheduler returned {len(results)} results, expected {B * G}"
            )
        for token_ids, logprobs in results:
            if len(token_ids) != len(logprobs):
                raise RuntimeError(
                    "Rollout scheduler returned misaligned token IDs and logprobs"
                )
        # Each element is (token_ids, logprobs); pad to max length.
        max_len = 0
        for token_ids, _lp in results:
            max_len = max(max_len, len(token_ids))
        max_len = max(max_len, 1)
        device = self.scheduler.device
        P_len = max(len(ids) for ids in flat_prompt_ids)
        prompts_tensor = torch.zeros(B, P_len, dtype=torch.long, device=device)
        prompt_mask = torch.zeros(B, P_len, dtype=torch.bool, device=device)
        for i, ids in enumerate(flat_prompt_ids):
            prompts_tensor[i, -len(ids) :] = torch.tensor(
                ids, dtype=torch.long, device=device
            )
            prompt_mask[i, -len(ids) :] = True
        responses = torch.full((B, G, max_len), _PAD, dtype=torch.long, device=device)
        response_mask = torch.zeros((B, G, max_len), dtype=torch.bool, device=device)
        logprobs_old = torch.zeros((B, G, max_len), dtype=torch.float, device=device)
        flat_idx = 0
        response_texts: List[List[str]] = [[] for _ in range(B)]
        for i in range(B):
            for g in range(G):
                token_ids, lps = results[flat_idx]
                flat_idx += 1
                n = len(token_ids)
                if n:
                    responses[i, g, :n] = torch.tensor(
                        token_ids, dtype=torch.long, device=device
                    )
                    response_mask[i, g, :n] = True
                    logprobs_old[i, g, :n] = torch.tensor(
                        lps, dtype=torch.float, device=device
                    )
                response_texts[i].append(
                    self.tokenizer.decode(token_ids, skip_special_tokens=True)
                )
        return RawRollout(
            prompts=prompts_tensor,
            prompt_mask=prompt_mask,
            responses=responses,
            response_mask=response_mask,
            logprobs_old=logprobs_old,
            prompt_texts=prompt_texts,
            response_texts=response_texts,
        )
    def _prepare_prompts(self, batch: Dict) -> Tuple[List[str], List[List[int]]]:
        """Render batch prompts to ``(texts, token_id_lists)``.
        Returns two parallel lists of length B (number of prompts in
        the batch).  Dispatches by batch keys:
        - ``"messages"``: treated as a pre-built message list per sample.
        - ``"instruction"`` (optionally ``"input"`` and ``"output"``): mapped
          to ``system`` / ``user`` / ``assistant`` messages respectively.
        Both paths go through the tokenizer's chat template with
        ``add_generation_prompt=True``.
        """
        if "messages" in batch:
            messages_list = batch["messages"]
        elif "instruction" in batch:
            instructions = batch["instruction"]
            B = len(instructions)
            inputs = batch.get("input") or [""] * B
            outputs = batch.get("output") or [""] * B
            messages_list = [
                self._instruction_to_messages(i, u, o)
                for i, u, o in zip(instructions, inputs, outputs)
            ]
        else:
            raise ValueError(
                "Rollout batch must contain either 'messages' or "
                "'instruction' (optionally 'input'/'output'); got keys: "
                f"{list(batch.keys())}"
            )
        try:
            prompt_texts = self.tokenizer.apply_chat_template(
                messages_list, tokenize=False, add_generation_prompt=True
            )
            if (
                not isinstance(prompt_texts, list)
                or len(prompt_texts) != len(messages_list)
                or not all(isinstance(text, str) for text in prompt_texts)
            ):
                raise TypeError("Tokenizer does not support batched chat templates")
            flat_prompt_ids = self.tokenizer.encode(prompt_texts)
            if len(flat_prompt_ids) != len(messages_list) or not all(
                isinstance(ids, list) for ids in flat_prompt_ids
            ):
                raise TypeError("Tokenizer does not support batched encoding")
        except (TypeError, IndexError, KeyError):
            # Keep compatibility with lightweight tokenizer adapters that only
            # implement the single-conversation template API.
            prompt_texts = []
            flat_prompt_ids = []
            for messages in messages_list:
                text = self.tokenizer.apply_chat_template(
                    messages, tokenize=False, add_generation_prompt=True
                )
                ids = self.tokenizer.apply_chat_template(
                    messages, tokenize=True, add_generation_prompt=True
                )
                prompt_texts.append(text)
                flat_prompt_ids.append(list(ids))
        return prompt_texts, flat_prompt_ids
    @staticmethod
    def _instruction_to_messages(
        instruction: str, inp: str = "", output: str = ""
    ) -> List[Dict[str, str]]:
        """Map instruction/input/output to chat messages.
        Role mapping follows the convention used throughout the
        preprocessing pipeline: ``instruction`` → system, ``input`` →
        user, ``output`` → assistant.  Empty fields are skipped so a
        bare instruction produces a ``[system]`` list and the chat
        template's ``add_generation_prompt`` adds the assistant header
        for sampling.
        """
        messages: List[Dict[str, str]] = []
        if instruction:
            messages.append({"role": "system", "content": instruction})
        if inp:
            messages.append({"role": "user", "content": inp})
        if output:
            messages.append({"role": "assistant", "content": output})
        return messages
 class RolloutRunner:
    """Produces :class:`RolloutResult` from a prompt batch.
    Composes a :class:`RolloutGenerator` (generation + decoding) with a
    :class:`BaseRewardModel` (scoring).  Maintains an internal cache so
    the same batch prompt can be replayed for multiple gradient steps.
    A new rollout is triggered every ``rollout_interval`` calls to
    :meth:`step` (or after :meth:`clear_cache`).
    The ``__call__`` contract returns a ``(RolloutResult, is_fresh)``
    tuple — callers must use the boolean to detect a refreshed rollout
    rather than relying on object identity.
    Usage::
        generator = RolloutGenerator(policy, tokenizer, pipeline, ...)
        runner = RolloutRunner(generator, reward_model, rollout_interval=512)
        result, is_fresh = runner(prompt_batch)
        if is_fresh:
            ...  # e.g. sync behaviour policy
    """
    def __init__(
        self,
        generator: RolloutGenerator,
        reward_model: BaseRewardModel,
        rollout_interval: int = 512,
    ):
        self.generator = generator
        self.reward_model = reward_model
        self.rollout_interval = rollout_interval
        self._cache: Optional[RolloutResult] = None
        self._cache_key = None
        self._steps_since_rollout: int = 0
    def step(self):
        """Advance the internal counter (call once per optimizer step)."""
        self._steps_since_rollout += 1
    def clear_cache(self):
        """Force next call to re-run rollout."""
        self._cache = None
        self._cache_key = None
    @staticmethod
    def _batch_key(batch: Dict):
        """Build a stable key for the prompt fields accepted by the generator."""
        def freeze(value):
            if isinstance(value, dict):
                return tuple(sorted((key, freeze(val)) for key, val in value.items()))
            if isinstance(value, (list, tuple)):
                return tuple(freeze(item) for item in value)
            return value
        fields = ("messages", "instruction", "input", "output")
        return tuple(
            (field, freeze(batch[field])) for field in fields if field in batch
        )
    def _score(self, raw: RawRollout) -> RolloutResult:
        rewards = self.reward_model.score(raw.prompt_texts, raw.response_texts)
        if not isinstance(rewards, Tensor):
            rewards = torch.as_tensor(rewards, dtype=torch.float32)
        expected_shape = raw.responses.shape[:2]
        if rewards.shape != expected_shape:
            raise ValueError(
                f"Reward model returned shape {tuple(rewards.shape)}, "
                f"expected {tuple(expected_shape)}"
            )
        if not torch.isfinite(rewards).all():
            raise ValueError("Reward model returned non-finite values")
        device = raw.prompts.device
        return RolloutResult(
            prompts=raw.prompts,
            prompt_mask=raw.prompt_mask,
            responses=raw.responses,
            response_mask=raw.response_mask,
            rewards=rewards.to(device=device),
            logprobs_old=raw.logprobs_old,
            prompt_texts=raw.prompt_texts,
            response_texts=raw.response_texts,
        )
    def __call__(self, batch: Dict[str, Tensor]) -> Tuple[RolloutResult, bool]:
        """Return ``(cached or fresh) RolloutResult`` plus an ``is_fresh`` flag.
        Triggers a new rollout when ``_steps_since_rollout >= rollout_interval``
        or when the cache is empty.
        """
        cache_key = self._batch_key(batch)
        if (
            self._cache is None
            or cache_key != self._cache_key
            or self._steps_since_rollout >= self.rollout_interval
        ):
            raw = self.generator.generate(batch)
            self._cache = self._score(raw)
            self._cache_key = cache_key
            self._steps_since_rollout = 0
            return self._cache, True
        return self._cache, False
@@ -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]):
        """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,57 +1,49 @@
 """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, List, Optional, TypedDict, Union
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from torch import Tensor
 from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
 from torch.nn.parallel import DistributedDataParallel as DDP
 from astrai.factory import BaseFactory
 from astrai.model.components.mlp import RouterStats
 from astrai.parallel.executor import broadcast_state_dict
 from astrai.trainer.rollout import RolloutResult
-def unwrap_model(model: nn.Module) -> nn.Module:
+class LossOutput(TypedDict):
-    if isinstance(model, DDP):
+    loss: Tensor
-        return model.module
+    metrics: Dict[str, float]
    if isinstance(model, FSDP):
        return model._fsdp_wrapped_module
    return model
-def create_ref_model(model: nn.Module) -> nn.Module:
+class LogprobsOutput(TypedDict):
-    """Create a reference model for DPO/GRPO training.
+    logprobs: Tensor
-
+    aux_loss: Optional[Tensor]
-    Handles DDP-wrapped models safely by unwrapping first,
+    router_stats: Optional[List[RouterStats]]
    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()}
 def get_logprobs(
-    model: Union[nn.Module, Callable[..., Dict[str, Tensor]]],
+    model: nn.Module,
    input_ids: Tensor,
-    mask: Tensor,
+    attn_mask: Tensor,
    loss_mask: Tensor,
    reduction: str,
-):
+) -> LogprobsOutput:
    """Compute token-wise log probabilities from model outputs.
    Args:
        model: The language model
        input_ids: Input token IDs of shape [batch_size, seq_len]
-        mask: Attention mask of shape [batch_size, seq_len]
+        attn_mask: Attention mask passed to the model (may include causal).
        loss_mask: Per-token mask for loss reduction.
        reduction: How to reduce over sequence dimension ("mean", "sum", "none")
    Returns:
@@ -64,9 +56,13 @@ def get_logprobs(
        )
    shifted_input_ids = input_ids[:, 1:]
-    shifted_mask = mask[:, 1:]
+    shifted_loss_mask = loss_mask[:, 1:]
-    logits = model(input_ids[:, :-1], mask[:, :-1])["logits"]
+    outputs = model(
        input_ids[:, :-1],
        attn_mask[:, :, :-1, :-1] if attn_mask.dim() == 4 else attn_mask[:, :-1],
    )
    logits = outputs["logits"]
    log_probs = torch.log_softmax(logits.float(), dim=-1)
    token_logprobs = torch.gather(
@@ -74,24 +70,122 @@ def get_logprobs(
    ).squeeze(-1)
    if reduction == "mean":
-        return (token_logprobs * shifted_mask).sum(dim=-1) / shifted_mask.sum(
+        logprobs = (token_logprobs * shifted_loss_mask).sum(
            dim=-1
-        ).clamp(min=1.0)
+        ) / shifted_loss_mask.sum(dim=-1).clamp(min=1.0)
    elif reduction == "sum":
-        return (token_logprobs * shifted_mask).sum(dim=-1)
+        logprobs = (token_logprobs * shifted_loss_mask).sum(dim=-1)
    else:
-        return token_logprobs * shifted_mask
+        logprobs = token_logprobs * shifted_loss_mask
    return {
        "logprobs": logprobs,
        "aux_loss": outputs.get("aux_loss"),
        "router_stats": outputs.get("router_stats"),
    }
 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)
 def _collect_moe_diagnostics(
    router_stats_list: List[RouterStats],
 ) -> Dict[str, float]:
    """Collect MoE routing diagnostic metrics from per-layer router stats.
    Args:
        router_stats_list: One :class:`RouterStats` dict per MoE layer with
            keys ``probs`` (N, E) and ``topk_indices`` (N, K), both detached.
    Returns:
        Dict with keys: router_entropy, dead_expert_fraction,
        load_imbalance_mean, load_imbalance_max.  Values are averaged
        across layers.
    """
    layer_entropies: List[Tensor] = []
    layer_dead_fractions: List[Tensor] = []
    layer_imbalance_means: List[Tensor] = []
    layer_imbalance_maxs: List[Tensor] = []
    for stats in router_stats_list:
        probs = stats["probs"].float()
        topk_indices = stats["topk_indices"]
        num_experts = probs.shape[-1]
        if num_experts == 0:
            continue
        probs = probs.reshape(-1, num_experts)
        if probs.numel() == 0:
            continue
        # Router entropy
        entropy = -(probs * torch.log(probs.clamp_min(1e-8))).sum(dim=-1).mean()
        # Load from the actual dispatch: one-hot sum of top-k assignments.
        expert_counts = F.one_hot(topk_indices, num_experts).sum(dim=(0, 1)).float()
        ideal_load = expert_counts.mean()  # N*K / E
        load_ratios = expert_counts / max(float(ideal_load), 1.0)
        imbalance_mean = (load_ratios - 1.0).abs().mean()
        imbalance_max = load_ratios.max()
        dead_fraction = (expert_counts == 0).float().mean()
        layer_entropies.append(entropy)
        layer_dead_fractions.append(dead_fraction)
        layer_imbalance_means.append(imbalance_mean)
        layer_imbalance_maxs.append(imbalance_max)
    if not layer_entropies:
        return {}
    return {
        "router_entropy": float(torch.stack(layer_entropies).mean().cpu().item()),
        "dead_expert_fraction": float(
            torch.stack(layer_dead_fractions).mean().cpu().item()
        ),
        "load_imbalance_mean": float(
            torch.stack(layer_imbalance_means).mean().cpu().item()
        ),
        "load_imbalance_max": float(
            torch.stack(layer_imbalance_maxs).mean().cpu().item()
        ),
    }
 class BaseStrategy(ABC):
-    """Abstract base class for training strategies."""
+    """Abstract base class for training strategies.
    When a :class:`~astrai.trainer.rollout.RolloutRunner` is injected via
    :meth:`set_rollout_runner`, the strategy transparently switches to
    online mode: each ``__call__`` produces a :class:`RolloutResult`,
    converts it to a training batch via :meth:`prepare_from_rollout`, and
    then computes the loss.  Without a runner the strategy runs in
    offline mode and consumes the batch directly.
    """
    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.moe_aux_loss_coef = kwargs.pop("moe_aux_loss_coef", 0.01)
        self._moe_metrics: Dict[str, float] = {}
        self.extra_kwargs = kwargs
        self._rollout_runner = None
    @abstractmethod
    def compute_loss(self, batch: Dict[str, Tensor]) -> Tensor:
@@ -105,9 +199,96 @@ class BaseStrategy(ABC):
        """
        raise NotImplementedError
-    def __call__(self, batch: Dict[str, Tensor]) -> Tensor:
+    def compute_loss_output(self, batch: Dict[str, Tensor]) -> LossOutput:
-        """Allow calling strategy directly as a callable."""
+        return self._normalize_output(self.compute_loss(batch))
-        return self.compute_loss(batch)
+
    def _loss_output(
        self,
        task_loss: Tensor,
        metrics: Dict[str, Tensor],
        aux_loss: Optional[Tensor] = None,
        router_stats: Optional[List[RouterStats]] = None,
    ) -> LossOutput:
        total_loss = task_loss
        if aux_loss is not None:
            weighted_aux_loss = self.moe_aux_loss_coef * aux_loss
            total_loss = total_loss + weighted_aux_loss
            metrics["moe_aux_loss"] = aux_loss
            metrics["moe_aux_loss_weighted"] = weighted_aux_loss
            self._refresh_moe_diagnostics(aux_loss, router_stats)
        metrics["loss"] = total_loss
        return {
            "loss": total_loss,
            "metrics": {name: value.detach().item() for name, value in metrics.items()},
        }
    @staticmethod
    def _normalize_output(output: Union[LossOutput, Tensor]) -> LossOutput:
        if isinstance(output, dict):
            return output
        return {"loss": output, "metrics": {"loss": output.detach().item()}}
    def supports_online(self) -> bool:
        """Whether this strategy can operate with a rollout runner.
        Base implementation returns ``False``; strategies that implement
        :meth:`prepare_from_rollout` should override to return ``True``.
        """
        return False
    def set_rollout_runner(self, runner):
        """Inject a :class:`RolloutRunner` to enable online rollout mode."""
        self._rollout_runner = runner
    def prepare_from_rollout(self, result: RolloutResult) -> Dict[str, Tensor]:
        """Map a :class:`RolloutResult` to the batch layout expected by
        :meth:`compute_loss`.
        Strategies that return ``True`` from :meth:`supports_online` must
        override this.  Default raises :class:`NotImplementedError`.
        """
        raise NotImplementedError(
            f"{type(self).__name__} does not support online rollout"
        )
    def _on_rollout_refresh(self):
        """Hook fired when a fresh rollout result is produced.
        Override to refresh stale state (e.g. syncing the behaviour
        policy).  Default is a no-op.
        """
        pass
    def _refresh_moe_diagnostics(
        self,
        aux_loss: Tensor,
        router_stats: Optional[List[RouterStats]] = None,
    ) -> None:
        """Collect MoE routing diagnostics from the latest forward pass.
        Populates ``self._moe_metrics`` with router entropy, dead expert
        fraction, load imbalance, and aux_loss.  Called from
        :meth:`_loss_output` when an MoE aux loss is present.
        """
        self._moe_metrics = _collect_moe_diagnostics(router_stats or [])
        self._moe_metrics["aux_loss"] = float(aux_loss.detach().cpu().item())
    def on_optimizer_step(self):
        """Advance online rollout state after a successful optimizer step."""
        if self._rollout_runner is not None:
            self._rollout_runner.step()
    def __call__(self, batch: Dict[str, Tensor]) -> LossOutput:
        """Run offline or online forward depending on runner injection."""
        if self._rollout_runner is None:
            return self.compute_loss_output(batch)
        result, is_fresh = self._rollout_runner(batch)
        if is_fresh:
            self._on_rollout_refresh()
        train_batch = self.prepare_from_rollout(result)
        return self.compute_loss_output(train_batch)
 class StrategyFactory(BaseFactory["BaseStrategy"]):
@@ -124,32 +305,6 @@ class StrategyFactory(BaseFactory["BaseStrategy"]):
        strategy = StrategyFactory.create("custom", model, device)
    """
    @classmethod
    def _validate_component(cls, strategy_cls: type):
        """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,16 +315,27 @@ class SEQStrategy(BaseStrategy):
    """Standard next-token prediction training strategy.
    Computes cross-entropy loss for next token prediction.
    Optionally adds MoE load balancing auxiliary loss.
    """
-    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:
        return self.compute_loss_output(batch)["loss"]
    def compute_loss_output(self, batch: Dict[str, Tensor]) -> LossOutput:
        batch = move_to_device(batch, self.device)
        input_ids, target_ids = batch["input_ids"], batch["target_ids"]
-        logits = self.model(input_ids=input_ids)["logits"]
+        outputs = self.model(input_ids=input_ids)
        logits = outputs["logits"]
        loss = F.cross_entropy(
            input=logits.flatten(0, 1).float(),
@@ -177,7 +343,12 @@ class SEQStrategy(BaseStrategy):
            label_smoothing=self.label_smoothing,
        )
-        return loss
+        return self._loss_output(
            loss,
            {"task_loss": loss},
            outputs.get("aux_loss"),
            outputs.get("router_stats"),
        )
@StrategyFactory.register("sft")
@@ -185,23 +356,38 @@ class SFTStrategy(BaseStrategy):
    """Supervised Fine-tuning strategy with loss masking.
    Applies cross-entropy loss only to tokens where loss_mask is True.
    Optionally adds MoE load balancing auxiliary loss.
    """
-    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:
        return self.compute_loss_output(batch)["loss"]
    def compute_loss_output(self, batch: Dict[str, Tensor]) -> LossOutput:
        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)
        outputs = self.model(
            input_ids=input_ids, position_ids=position_ids, input_mask=input_mask
        )
        logits = outputs["logits"]
        loss = F.cross_entropy(
            input=logits.flatten(0, 1).float(),
@@ -210,7 +396,12 @@ class SFTStrategy(BaseStrategy):
            label_smoothing=self.label_smoothing,
        )
-        return loss
+        return self._loss_output(
            loss,
            {"task_loss": loss},
            outputs.get("aux_loss"),
            outputs.get("router_stats"),
        )
@StrategyFactory.register("dpo")
@@ -225,29 +416,54 @@ 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
    def compute_loss(self, batch: Dict[str, Tensor]) -> Tensor:
        return self.compute_loss_output(batch)["loss"]
    def compute_loss_output(self, batch: Dict[str, Tensor]) -> LossOutput:
        batch = move_to_device(batch, self.device)
        chosen_ids, rejected_ids = batch["chosen"], batch["rejected"]
        chosen_mask, rejected_mask = batch["chosen_mask"], batch["rejected_mask"]
        concat_ids = torch.cat([chosen_ids, rejected_ids], dim=0)
-        concat_mask = torch.cat([chosen_mask, rejected_mask], dim=0)
+        concat_loss_mask = torch.cat([chosen_mask, rejected_mask], dim=0)
-        log_pi = get_logprobs(self.model, concat_ids, concat_mask, self.reduction)
+        # Build full attention mask: key-padding + causal
        key_pad = concat_ids.bool()[:, None, None, :]  # [B*2, 1, 1, S]
        S = key_pad.shape[-1]
        causal = torch.tril(
            torch.ones(S, S, dtype=torch.bool, device=concat_ids.device)
        )[None, None, :, :]  # [1, 1, S, S]
        full_mask = key_pad & causal  # [B*2, 1, S, S] — composed
        policy_output = get_logprobs(
            self.model,
            concat_ids,
            full_mask,
            concat_loss_mask,
            self.reduction,
        )
        log_pi = policy_output["logprobs"]
        aux_loss = policy_output["aux_loss"]
        with torch.no_grad():
-            log_ref = get_logprobs(
+            ref_output = get_logprobs(
-                self.ref_model, concat_ids, concat_mask, self.reduction
+                self.ref_model,
                concat_ids,
                full_mask,
                concat_loss_mask,
                self.reduction,
            )
            log_ref = ref_output["logprobs"]
        log_pi_chosen = log_pi[: chosen_ids.shape[0]]
        log_pi_rejected = log_pi[chosen_ids.shape[0] :]
@@ -260,48 +476,87 @@ class DPOStrategy(BaseStrategy):
        ratio_diff = pi_log_ratio - ref_log_ratio
        dpo_loss = -F.logsigmoid(self.beta * ratio_diff).mean()
-        return dpo_loss
+        return self._loss_output(
            dpo_loss,
            {"dpo_loss": dpo_loss},
            aux_loss,
            policy_output.get("router_stats"),
        )
    def supports_online(self) -> bool:
        return True
    def prepare_from_rollout(self, result: RolloutResult) -> Dict[str, Tensor]:
        """Pick best/worst response per prompt by reward as chosen/rejected."""
        rewards = result.rewards
        responses = result.responses
        masks = result.response_mask
        best = rewards.argmax(dim=-1)
        worst = rewards.argmin(dim=-1)
        B = responses.shape[0]
        idx = torch.arange(B, device=responses.device)
        chosen = responses[idx, best]
        chosen_mask = masks[idx, best].float()
        rejected = responses[idx, worst]
        rejected_mask = masks[idx, worst].float()
        return {
            "chosen": chosen,
            "chosen_mask": chosen_mask,
            "rejected": rejected,
            "rejected_mask": rejected_mask,
        }
@StrategyFactory.register("grpo")
 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()
+        state_dict = self.executor.unwrap_model(self.model)
-        self.ref_model.load_state_dict(ref_state)
+        if self.executor.use_distributed:
            state_dict = broadcast_state_dict(state_dict)
        if state_dict is not None:
            self.old_model.load_state_dict(state_dict)
    def compute_loss(self, batch: Dict[str, Tensor]) -> Tensor:
-        self._step += 1
+        return self.compute_loss_output(batch)["loss"]
        if self._step % self.sync_interval == 0:
            self.sync_ref_model()
    def compute_loss_output(self, batch: Dict[str, Tensor]) -> LossOutput:
        batch = move_to_device(batch, self.device)
        prompts = batch["prompts"]
        responses = batch["responses"]
@@ -312,33 +567,112 @@ 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_mask = batch.get("prompt_mask")
        if prompt_mask is None:
            prompt_mask = prompts.ne(0)
        prompt_mask_expanded = (
            prompt_mask.unsqueeze(1).expand(-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, dtype=torch.bool), masks_flat], dim=-1
        )
        log_probs_policy = log_probs_policy.view(batch_size, group_size)
        # Build full attention mask: key-padding + causal
        key_pad = torch.cat([prompt_mask_expanded, masks_flat.bool()], dim=-1)[
            :, None, None, :
        ]
        S = key_pad.shape[-1]
        causal = torch.tril(
            torch.ones(S, S, dtype=torch.bool, device=full_sequences.device)
        )[None, None, :, :]
        attn_mask = key_pad & causal
        # 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).
        policy_output = get_logprobs(
            self.model, full_sequences, attn_mask, full_masks, "none"
        )
        token_log_probs_policy = policy_output["logprobs"]
        aux_loss = policy_output["aux_loss"]
        token_log_probs_policy = token_log_probs_policy[:, prompt_len - 1 :]
        with torch.no_grad():
-            log_probs_ref = get_logprobs(
+            old_output = get_logprobs(
-                self.ref_model, full_sequences, full_masks, self.reduction
+                self.old_model, full_sequences, attn_mask, full_masks, "none"
            )
-            log_probs_ref = log_probs_ref.view(batch_size, group_size)
+            token_log_probs_old = old_output["logprobs"]
            token_log_probs_old = token_log_probs_old[:, prompt_len - 1 :]
            ref_output = get_logprobs(
                self.ref_model, full_sequences, attn_mask, full_masks, "none"
            )
            token_log_probs_ref = ref_output["logprobs"]
            token_log_probs_ref = token_log_probs_ref[:, 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
-        policy_loss = -torch.min(surr1, surr2).mean()
+        # KL penalty to frozen reference model with k1 estimator (non-negative):
-        kl_penalty = self.kl_coef * (log_probs_policy - log_probs_ref).square().mean()
+        # k1 = π_ref / π_θ - log(π_ref / π_θ) - 1, where π_ref / π_θ = exp(log_ref - log_policy).
-        total_loss = policy_loss + kl_penalty
+        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
-        return total_loss
+        task_loss = policy_loss + kl_penalty
        return self._loss_output(
            task_loss,
            {"policy_loss": policy_loss, "kl_loss": kl_penalty},
            aux_loss,
            policy_output.get("router_stats"),
        )
    def supports_online(self) -> bool:
        return True
    def prepare_from_rollout(self, result: RolloutResult) -> Dict[str, Tensor]:
        return {
            "prompts": result.prompts,
            "prompt_mask": result.prompt_mask,
            "responses": result.responses,
            "masks": result.response_mask,
            "rewards": result.rewards,
        }
    def _on_rollout_refresh(self):
        """Sync the behaviour policy whenever a fresh rollout arrives."""
        self.sync_old_model()
 # Factory aliases: online variants use the same strategy class; the
 # ``RolloutRunner`` is injected by ``TrainContextBuilder`` to enable
 # online mode, so no separate subclass is needed.
 StrategyFactory.register("online_grpo")(GRPOStrategy)
 StrategyFactory.register("online_dpo")(DPOStrategy)
@@ -9,23 +9,23 @@ from typing import IO, Callable, List, Optional, Protocol, runtime_checkable
 import torch
 import torch.distributed as dist
 import torch.nn as nn
 from torch.nn.utils import clip_grad_norm_
 from torch.utils.checkpoint import checkpoint as torch_checkpoint
 from tqdm import tqdm
 from astrai.factory import BaseFactory
 from astrai.parallel import only_on_rank
-from astrai.parallel.setup import get_current_device, get_rank
+from astrai.parallel.setup import get_current_device
 from astrai.serialization import Checkpoint
 from astrai.trainer.metric_util import (
-    ctx_get_grad_max,
+    ctx_get_dead_expert_fraction,
    ctx_get_grad_mean,
    ctx_get_grad_min,
    ctx_get_grad_nan_num,
    ctx_get_grad_norm,
-    ctx_get_grad_std,
+    ctx_get_grad_snr,
    ctx_get_load_imbalance_max,
    ctx_get_load_imbalance_mean,
    ctx_get_loss,
    ctx_get_lr,
    ctx_get_moe_aux_loss,
    ctx_get_router_entropy,
    ctx_get_val_loss,
 )
 from astrai.trainer.train_context import TrainContext
@@ -86,7 +86,9 @@ class GradientClippingCallback(TrainCallback):
        self.max_grad_norm = max_grad_norm
    def on_optimizer_step(self, context: TrainContext):
-        clip_grad_norm_(context.model.parameters(), self.max_grad_norm)
+        context.grad_norm = context.executor.clip_grad_norm(
            context.model, self.max_grad_norm
        )
@CallbackFactory.register("gradient_checkpointing")
@@ -143,33 +145,38 @@ class CheckpointCallback(TrainCallback):
        self.interval = interval
        self.weight_only = weight_only
        self.save_extra_fn = save_extra_fn or CheckpointCallback.save_extra
-        self.last_ckpt_iter = 0
+        self.last_ckpt_step = None
    def on_train_begin(self, context: TrainContext):
        self.last_ckpt_step = context.optimizer_step
    def _save_checkpoint(self, context: TrainContext):
-        unwrapped = context.executor.unwrap_model(context.model)
+        self.last_ckpt_step = context.optimizer_step
        state_dict = unwrapped.state_dict()
        self.last_ckpt_iter = context.iteration
-        if get_rank() == 0:
+        with context.executor.checkpoint_context(context.model) as state_dict:
-            save_path = os.path.join(
+            if state_dict is not None:
-                self.save_dir, f"epoch_{context.epoch}_iter_{context.iteration}"
+                save_path = os.path.join(
-            )
+                    self.save_dir,
-            extra = self.save_extra_fn(context)
+                    f"epoch_{context.epoch}_step_{context.optimizer_step}",
-            context.checkpoint = Checkpoint(
+                )
-                state_dict=state_dict,
+                extra = self.save_extra_fn(context)
-                epoch=context.epoch,
+                meta = context.config.to_dict()
-                iteration=context.iteration,
+                context.checkpoint = Checkpoint(
-                extra=extra,
+                    state_dict=state_dict,
-                config=context.model_config,
+                    epoch=context.epoch,
-            )
+                    consumed_samples=context.consumed_samples,
-            context.checkpoint.save(save_path)
+                    config=context.model_config,
                    extra=extra,
                    meta=meta,
                )
                context.checkpoint.save(save_path)
    def on_batch_end(self, context: TrainContext):
-        if context.iteration - self.last_ckpt_iter >= self.interval:
+        if context.optimizer_step - self.last_ckpt_step >= self.interval:
            self._save_checkpoint(context)
    def on_train_end(self, context: TrainContext):
-        if context.iteration != self.last_ckpt_iter:
+        if context.optimizer_step != self.last_ckpt_step:
            self._save_checkpoint(context)
    def on_error(self, context: TrainContext):
@@ -201,20 +208,24 @@ class ProgressBarCallback(TrainCallback):
    @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):
        postfix = {
            "step": f"{context.optimizer_step:d}",
            "loss": f"{context.loss:.4f}",
            "lr": f"{context.optimizer.param_groups[-1]['lr']:.2e}",
        }
-        if context.val_loss > 0:
+        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)
@@ -226,22 +237,22 @@ 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,
+        ckpt_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.ckpt_dir = Path(ckpt_dir) if ckpt_dir else Path.cwd() / "checkpoint"
        self.log_dir.mkdir(parents=True, exist_ok=True)
        self.log_cache = []
@@ -250,53 +261,51 @@ class MetricLoggerCallback(TrainCallback):
            "lr": ctx_get_lr,
            "val_loss": ctx_get_val_loss,
            "grad_norm": ctx_get_grad_norm,
-            "grad_std": ctx_get_grad_std,
+            "grad_snr": ctx_get_grad_snr,
-            "grad_max": ctx_get_grad_max,
+            "moe_aux_loss": ctx_get_moe_aux_loss,
-            "grad_min": ctx_get_grad_min,
+            "router_entropy": ctx_get_router_entropy,
-            "grad_mean": ctx_get_grad_mean,
+            "dead_expert_fraction": ctx_get_dead_expert_fraction,
-            "grad_nan_num": ctx_get_grad_nan_num,
+            "load_imbalance_mean": ctx_get_load_imbalance_mean,
            "load_imbalance_max": ctx_get_load_imbalance_max,
        }
-    def _get_log_data(self, context: TrainContext):
+    def _metrics(self, context: TrainContext, names):
-        return {
+        metrics = dict(context.metrics)
        for name in names:
            metric_fn = self._metric_funcs.get(name)
            if metric_fn is None:
                continue
            value = metric_fn(context)
            if value is not None:
                metrics[name] = value
        selected = set(context.metrics) | set(names)
        selected.discard("*")
        result = {name: metrics[name] for name in selected if name in metrics}
        if context.world_size > 1 and dist.is_initialized() and result:
            metric_names = sorted(result)
            values = torch.tensor(
                [result[name] for name in metric_names],
                dtype=torch.float32,
                device=get_current_device(),
            )
            dist.all_reduce(values, op=dist.ReduceOp.SUM)
            values /= context.world_size
            result.update(zip(metric_names, values.tolist()))
        return result
    @only_on_rank(0)
    def _append(self, event_type: str, context: TrainContext, **extra):
        entry = {
            "type": event_type,
            "timestamp": time.strftime("%Y-%m-%dT%H:%M:%S"),
            "epoch": context.epoch,
-            "iter": context.iteration,
+            "step": context.optimizer_step,
-            **{m: self._metric_funcs[m](context) for m in self.metrics},
+            "consumed_samples": context.consumed_samples,
            **extra,
        }
        self.log_cache.append(entry)
-    @only_on_rank(0)
+    def _run_validation(self, context: TrainContext) -> float:
    def _add_log(self, log_data):
        self.log_cache.append(log_data)
    @only_on_rank(0)
    def _save_log(self, epoch, iter):
        log_file = self.log_dir / f"epoch_{epoch}_iter_{iter}_metric.jsonl"
        with open(log_file, "w") as f:
            for log in self.log_cache:
                f.write(json.dumps(log) + "\n")
    def on_batch_end(self, context):
        if context.iteration % self.log_interval == 0:
            log_data = self._get_log_data(context)
            self._add_log(log_data)
        if context.iteration - self.last_log_iter >= self.save_interval:
            self._save_log(context.epoch, context.iteration)
            self.last_log_iter = context.iteration
    def on_train_end(self, context):
        if context.iteration != self.last_log_iter:
            self._save_log(context.epoch, context.iteration)
    def on_error(self, context):
        self._save_log(context.epoch, context.iteration)
@CallbackFactory.register("validation")
 class ValidationCallback(TrainCallback):
    def _run_validation(self, context: TrainContext):
        context.model.eval()
        total_loss = 0.0
@@ -304,31 +313,62 @@ class ValidationCallback(TrainCallback):
        with torch.no_grad():
            for batch in context.val_dataloader:
-                loss = context.strategy(batch)
+                loss_output = context.strategy(batch)
-                total_loss += loss.item()
+                total_loss += loss_output["loss"].item()
                num_batches += 1
        avg_loss = total_loss / max(num_batches, 1)
        if context.world_size > 1 and dist.is_initialized():
-            loss_tensor = torch.tensor([avg_loss], device=get_current_device())
+            stats = torch.tensor(
-            dist.all_reduce(loss_tensor, op=dist.ReduceOp.AVG)
+                [total_loss, float(num_batches)], device=get_current_device()
-            avg_loss = loss_tensor.item()
+            )
            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.val_loss = avg_loss
        context.model.train()
        return avg_loss
-        step_count = context.iteration // context.config.grad_accum_steps
+    def on_train_begin(self, context: TrainContext):
-        logger.info(
+        self.last_log_flush_step = context.optimizer_step
            f"Epoch {context.epoch + 1}, Step {step_count}, Val Loss: {avg_loss:.4f}"
        )
-    def on_optimizer_step(self, context: TrainContext):
+    @only_on_rank(0)
-        if context.val_dataloader is None:
+    def _flush(self, epoch, step):
-            return
+        log_file = self.ckpt_dir / f"epoch_{epoch}_step_{step}" / "metric.jsonl"
-        cfg = context.config
+        log_file.parent.mkdir(parents=True, exist_ok=True)
-        if cfg.val_step <= 0:
+        with open(log_file, "w") as f:
-            return
+            for log in self.log_cache:
-        step_count = context.iteration // cfg.grad_accum_steps
+                f.write(json.dumps(log) + "\n")
-        if step_count % cfg.val_step == 0:
+
-            self._run_validation(context)
+    def on_optimizer_step(self, context):
        context.grad_snr_tracker.update(context.model)
        if (
            context.val_dataloader is not None
            and self.val_step > 0
            and context.optimizer_step >= self._next_val_step
        ):
            context.val_loss = self._run_validation(context)
            self._next_val_step = context.optimizer_step + self.val_step
            self._append("validation", context, val_loss=context.val_loss)
        step_metrics = [m for m in self.metrics if m != "val_loss"]
        self._append("step", context, **self._metrics(context, step_metrics))
        if context.optimizer_step - self.last_log_flush_step >= self.save_interval:
            self._flush(context.epoch, context.optimizer_step)
            self.last_log_flush_step = context.optimizer_step
    def on_epoch_end(self, context):
        self._append("epoch", context)
    def on_train_end(self, context):
        if (
            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._flush(context.epoch, context.optimizer_step)
@@ -1,19 +1,28 @@
 import logging
 import threading
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Optional, Self
+from typing import Any, Dict, Optional, Self
 import torch
 import torch.nn as nn
-from torch.utils.data import DataLoader
+from torch.utils.data import DataLoader, random_split
 from astrai.config.train_config import TrainConfig
-from astrai.dataset import ResumableDistributedSampler
+from astrai.dataset import RDSampler
 from astrai.inference.core.scheduler import InferenceScheduler
 from astrai.model.components.lora import inject_lora
-from astrai.parallel.executor import BaseExecutor, ExecutorFactory
+from astrai.parallel.executor import BaseExecutor, ExecutorFactory, create_ref_model
 from astrai.parallel.setup import get_current_device, get_rank, get_world_size
 from astrai.protocols import OptimizerProtocol, SchedulerProtocol
-from astrai.serialization import Checkpoint, load_json, load_model_weights
+from astrai.serialization import Checkpoint, load_json
 from astrai.tokenize import AutoTokenizer
 from astrai.trainer.metric_util import GradSNRTracker
 from astrai.trainer.rollout import RolloutGenerator, RolloutRunner
 from astrai.trainer.strategy import BaseStrategy, StrategyFactory
 logger = logging.getLogger(__name__)
@dataclass
 class TrainContext:
@@ -26,16 +35,35 @@ class TrainContext:
    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)
-    val_dataloader: DataLoader = field(default=None)
+    metrics: Dict[str, float] = field(default_factory=dict)
-    val_loss: float = field(default=0.0)
+    grad_norm: Optional[float] = field(default=None)
    grad_snr_tracker: GradSNRTracker = field(default_factory=GradSNRTracker)
    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)
    _stop_event: threading.Event = field(default_factory=threading.Event)
    @property
    def stop_requested(self) -> bool:
        return self._stop_event.is_set()
    def request_stop(self) -> None:
        self._stop_event.set()
    @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:
@@ -44,10 +72,12 @@ class TrainContextBuilder:
        config: TrainConfig,
    ):
        self.config = config
-        self._resume_dir: Optional[str] = None
+        self._param_path: Optional[str] = None
        self._resume: bool = False
-    def with_resume_dir(self, resume_dir: Optional[str]) -> Self:
+    def with_param_path(self, param_path: Optional[str], resume: bool = False) -> Self:
-        self._resume_dir = resume_dir
+        self._param_path = param_path
        self._resume = resume
        return self
    def build(self) -> TrainContext:
@@ -60,96 +90,129 @@ class TrainContextBuilder:
            **cfg.executor_kwargs,
        )
        model = cfg.model_fn()
        model = model.to(device=device)
        model_config = {}
-        if self._resume_dir:
+        if self._param_path:
-            config_path = Path(self._resume_dir) / "config.json"
+            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"):
+        preloaded_state_dict = None
-            model_config = model.config.to_dict()
+        preloaded_epoch = cfg.start_epoch
        preloaded_consumed = cfg.start_samples * get_world_size()
        preloaded_checkpoint = None
        if self._param_path:
            checkpoint = Checkpoint.load_any(self._param_path)
            if checkpoint is not None:
                preloaded_state_dict = checkpoint.state_dict
                if checkpoint.config:
                    model_config = checkpoint.config
                if self._resume:
                    preloaded_epoch = checkpoint.epoch
                    per_step = (
                        cfg.batch_per_device * get_world_size() * cfg.grad_accum_steps
                    )
                    preloaded_consumed = (
                        checkpoint.consumed_samples // per_step
                    ) * per_step
                    preloaded_checkpoint = checkpoint
        if not model_config and hasattr(cfg.model_fn(), "config"):
            model_config = cfg.model_fn().config.to_dict()
        def _before_wrap(m):
            m = m.to(device=device)
            if cfg.lora is not None:
                inject_lora(
                    m,
                    r=cfg.lora.r,
                    alpha=cfg.lora.alpha,
                    target_modules=set(cfg.lora.target_modules),
                )
            if preloaded_state_dict is not None:
                m.load_state_dict(preloaded_state_dict, strict=False)
            return m
        def _after_wrap(m):
            if cfg.compile_mode is not None:
                logger.info("torch.compile enabled (mode=%s)", cfg.compile_mode)
                m = torch.compile(m, mode=cfg.compile_mode)
            return m
        context = TrainContext(
            model=model,
            world_size=get_world_size(),
            rank=get_rank(),
            config=cfg,
            model_config=model_config,
            executor=executor,
            epoch=preloaded_epoch,
            consumed_samples=preloaded_consumed,
            checkpoint=preloaded_checkpoint,
        )
-        if self._resume_dir is not None:
+        context.model, context.optimizer, context.scheduler = executor.prepare(
-            resume_path = Path(self._resume_dir)
+            cfg.model_fn,
-            if (resume_path / "meta.json").exists():
+            cfg.optimizer_fn,
-                checkpoint = Checkpoint.load(self._resume_dir)
+            cfg.scheduler_fn,
-                state_dict = checkpoint.state_dict
+            before_wrap=_before_wrap,
-                if checkpoint.config:
+            after_wrap=_after_wrap,
-                    context.model_config = checkpoint.config
+        )
            else:
                checkpoint = None
                state_dict = load_model_weights(self._resume_dir)
            model.load_state_dict(state_dict, strict=False)
            if checkpoint is not None:
                context.epoch = cfg.start_epoch
                context.iteration = cfg.start_batch
            context.checkpoint = checkpoint
-        if cfg.lora is not None:
+        train_dataset = cfg.dataset
-            inject_lora(
+        val_dataset = cfg.val_dataset
-                model,
+
-                r=cfg.lora.r,
+        if val_dataset is None and cfg.val_split is not None:
-                alpha=cfg.lora.alpha,
+            n_total = len(cfg.dataset)
-                target_modules=set(cfg.lora.target_modules),
+            n_val = max(1, int(n_total * cfg.val_split))
            n_train = n_total - n_val
            generator = torch.Generator().manual_seed(cfg.random_seed)
            train_dataset, val_dataset = random_split(
                cfg.dataset, [n_train, n_val], generator=generator
            )
-        context.optimizer = cfg.optimizer_fn(model)
+        sampler_offset = context.consumed_samples // context.world_size
        context.scheduler = cfg.scheduler_fn(context.optimizer)
-        sampler_offset = context.iteration * cfg.batch_per_device
+        if self._resume and sampler_offset > 0:
-        sampler = ResumableDistributedSampler(
+            offset = context.world_size - 1
-            data_source=cfg.dataset,
+            num_samples_per_replica = (
                len(train_dataset) + offset
            ) // context.world_size
            if num_samples_per_replica > 0:
                context.epoch = sampler_offset // num_samples_per_replica
        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_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 cfg.val_dataset is not None:
+        if val_dataset is not None:
-            val_sampler = ResumableDistributedSampler(
+            val_sampler = RDSampler(
-                data_source=cfg.val_dataset,
+                data_source=val_dataset,
                start_epoch=0,
                start_iter=0,
                seed=cfg.random_seed,
                shuffle=False,
            )
            context.val_dataloader = DataLoader(
-                cfg.val_dataset,
+                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"):
@@ -158,11 +221,73 @@ class TrainContextBuilder:
                    if obj is not None:
                        obj.load_state_dict(extra[name])
        strategy_kwargs = dict(cfg.extra_kwargs)
        strategy_kwargs.setdefault("moe_aux_loss_coef", cfg.moe_aux_loss_coef)
        needs_ref = cfg.strategy in (
            "dpo",
            "grpo",
            "online_grpo",
            "online_dpo",
        )
        needs_old = cfg.strategy in ("grpo", "online_grpo")
        if needs_ref:
            strategy_kwargs["ref_model"] = create_ref_model(
                cfg.model_fn, executor=executor, model=context.model, device=device
            )
        if needs_old:
            strategy_kwargs["old_model"] = create_ref_model(
                cfg.model_fn, executor=executor, model=context.model, device=device
            )
        context.strategy = StrategyFactory.create(
            cfg.strategy,
            model=context.model,
            train_type=cfg.strategy,
            device=device,
-            **cfg.extra_kwargs,
+            executor=executor,
            **strategy_kwargs,
        )
        # Enable online rollout when the train_type is an ``online_*`` variant.
        is_online = cfg.strategy.startswith("online_")
        if is_online:
            if not context.strategy.supports_online():
                raise ValueError(
                    f"Strategy '{cfg.strategy}' does not support online rollout"
                )
            if cfg.reward_model_fn is None:
                raise ValueError("reward_model_fn is required for online RL strategies")
            tokenizer = AutoTokenizer.from_pretrained(self._param_path)
            reward_model = cfg.reward_model_fn()
            group_size = strategy_kwargs.get("group_size", 1)
            rollout_batch_size = group_size * max(1, cfg.batch_per_device)
            max_seq_len = getattr(context.model.config, "max_position_embeddings", None)
            scheduler = InferenceScheduler(
                model=context.model,
                tokenizer=tokenizer,
                max_batch_size=rollout_batch_size,
                max_seq_len=max_seq_len,
            )
            generator = RolloutGenerator(
                scheduler=scheduler,
                tokenizer=tokenizer,
                max_tokens=cfg.rollout_max_tokens,
                group_size=group_size,
                temperature=cfg.rollout_temperature,
                top_k=cfg.rollout_top_k,
                top_p=cfg.rollout_top_p,
            )
            runner = RolloutRunner(
                generator=generator,
                reward_model=reward_model,
                rollout_interval=cfg.rollout_interval,
            )
            context.strategy.set_rollout_runner(runner)
        return context
@@ -1,8 +1,14 @@
 import logging
 from typing import List, Optional
 import torch.distributed as dist
 from astrai.config import TrainConfig
 from astrai.parallel.setup import spawn_parallel_fn
 from astrai.signal_handler import (
    register_signal_handlers,
    unregister_signal_handlers,
 )
 from astrai.trainer.train_callback import (
    CallbackFactory,
    TrainCallback,
@@ -35,15 +41,14 @@ class Trainer:
                cfg.ckpt_interval,
            ),
            CallbackFactory.create(
-                "metric_logger",
+                "metric",
-                log_dir=cfg.log_dir,
+                ckpt_dir=cfg.ckpt_dir,
                save_interval=cfg.ckpt_interval,
                log_interval=cfg.log_interval,
                metrics=cfg.metrics,
                val_step=cfg.val_step,
            ),
            CallbackFactory.create("progress_bar", cfg.n_epoch),
            CallbackFactory.create("gradient_clipping", cfg.max_grad_norm),
            CallbackFactory.create("validation"),
        ]
        return callbacks
@@ -53,10 +58,13 @@ class Trainer:
            if method:
                method(context)
-    def _trainer_loop(self, resume_dir: Optional[str] = None):
+    def _trainer_loop(self, param_path: Optional[str] = None, resume: bool = False):
        context = (
-            TrainContextBuilder(self.train_config).with_resume_dir(resume_dir).build()
+            TrainContextBuilder(self.train_config)
            .with_param_path(param_path, resume=resume)
            .build()
        )
        register_signal_handlers(context)
        executor = context.executor
        self._call_callbacks("on_train_begin", context)
@@ -64,23 +72,30 @@ class Trainer:
            context.model.train()
            for epoch in range(context.epoch, context.config.n_epoch):
                if context.stop_requested:
                    break
                context.epoch = epoch
                self._call_callbacks("on_epoch_begin", context)
                for batch in context.dataloader:
-                    self._call_callbacks("on_batch_begin", context)
+                    if context.stop_requested:
-
+                        break
                    with executor.accumulate(context.model):
-                        loss = context.strategy(batch)
+                        self._call_callbacks("on_batch_begin", context)
-                        context.loss = loss.item()
+                        loss_output = context.strategy(batch)
-                        stand_loss = loss / executor.grad_accum_steps
+                        context.loss = loss_output["loss"].item()
                        context.metrics = loss_output["metrics"]
                        stand_loss = loss_output["loss"] / executor.grad_accum_steps
                        executor.backward(stand_loss)
-                        context.iteration += 1
+                        context.consumed_samples += (
                            context.config.batch_per_device * context.world_size
                        )
                        self._call_callbacks("on_batch_end", context)
                        if executor.sync_gradients:
                            self._call_callbacks("on_optimizer_step", context)
                            context.optimizer.step()
                            context.strategy.on_optimizer_step()
                            context.optimizer.zero_grad()
                            if context.scheduler:
@@ -88,14 +103,23 @@ class Trainer:
                self._call_callbacks("on_epoch_end", context)
            if context.stop_requested:
                logger.warning(
                    "Training interrupted by signal, saving emergency checkpoint..."
                )
                self._call_callbacks("on_error", context)
        except Exception as e:
            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)
            if executor.use_distributed and dist.is_initialized():
                dist.barrier()
            unregister_signal_handlers()
-    def train(self, resume_dir: Optional[str] = None):
+    def train(self, param_path: Optional[str] = None, resume: bool = False):
        cfg = self.train_config
        spawn_parallel_fn(
            self._trainer_loop,
@@ -105,5 +129,6 @@ class Trainer:
            master_port=cfg.master_port,
            device_type=cfg.device_type,
            start_method=cfg.start_method,
-            resume_dir=resume_dir,
+            param_path=param_path,
            resume=resume,
        )
@@ -0,0 +1,74 @@
 cmake_minimum_required(VERSION 3.18)
 project(astrai_kernels LANGUAGES CUDA CXX)
 set(CMAKE_CXX_STANDARD 17)
 set(CMAKE_CXX_STANDARD_REQUIRED ON)
 set(CMAKE_CUDA_STANDARD 17)
 find_package(CUDAToolkit REQUIRED)
 if(NOT DEFINED TORCH_HOME)
    set(TORCH_HOME "$ENV{TORCH_HOME}")
 endif()
 if(NOT TORCH_HOME)
    message(FATAL_ERROR "TORCH_HOME must point at the torch install dir (site-packages/torch)")
 endif()
 if(NOT DEFINED PYTHON_INCLUDE_DIR)
    set(PYTHON_INCLUDE_DIR "/usr/include/python${PYTHON_VERSION_MAJOR}.${PYTHON_VERSION_MINOR}")
 endif()
 if(NOT DEFINED ASTRAI_CUDA_ARCH)
    if(DEFINED ENV{ASTRAI_CUDA_ARCH})
        set(ASTRAI_CUDA_ARCH "$ENV{ASTRAI_CUDA_ARCH}")
    else()
        set(ASTRAI_CUDA_ARCH 80)
    endif()
 endif()
 set(TORCH_LIB_DIR "${TORCH_HOME}/lib")
 set(CUDA_LIB_DIR "/usr/local/cuda/lib64")
 set(CXX_FLAGS -O3 -funroll-loops)
 set(NVCC_FLAGS -O3
    --expt-relaxed-constexpr
    --use_fast_math
    "--ptxas-options=-O3,-v"
    --extra-device-vectorization
    --threads=16)
 set(TORCH_LIBS
    "${TORCH_LIB_DIR}/libtorch_python.so"
    "${TORCH_LIB_DIR}/libtorch_cuda.so"
    "${TORCH_LIB_DIR}/libc10_cuda.so"
    "${TORCH_LIB_DIR}/libtorch_cpu.so"
    "${TORCH_LIB_DIR}/libtorch.so"
    "${TORCH_LIB_DIR}/libc10.so"
    CUDA::cudart)
 set(CMAKE_CUDA_ARCHITECTURES "${ASTRAI_CUDA_ARCH}")
 set(KERNELS attn_decode attn_prefill attn_paged_decode attn_paged_prefill rotary_emb)
 foreach(name ${KERNELS})
    add_library(${name} MODULE "${CMAKE_CURRENT_SOURCE_DIR}/kernels/${name}.cu")
    target_compile_definitions(${name} PRIVATE TORCH_EXTENSION_NAME=${name})
    target_include_directories(${name} PRIVATE
        "${TORCH_HOME}/include"
        "${TORCH_HOME}/include/torch/csrc/api/include"
        "${PYTHON_INCLUDE_DIR}")
    target_link_libraries(${name} PRIVATE ${TORCH_LIBS})
    target_link_options(${name} PRIVATE "-Wl,-rpath,${TORCH_LIB_DIR}")
    target_compile_options(${name} PRIVATE
        $<$<COMPILE_LANGUAGE:CXX>:${CXX_FLAGS}>
        $<$<COMPILE_LANGUAGE:CUDA>:${NVCC_FLAGS}>)
    set_target_properties(${name} PROPERTIES
        PREFIX ""
        SUFFIX ".${PY_SOABI}.so"
        LIBRARY_OUTPUT_DIRECTORY "${CMAKE_CURRENT_SOURCE_DIR}/../astrai/extension/lib")
 endforeach()
@@ -0,0 +1,2 @@
 # Source directory for CUDA kernels — build-time only.
 # Compiled .so files live in astrAI/_ext/.
@@ -0,0 +1,66 @@
 #pragma once
 // Tensor layout for Q/K/V tensors passed to attention kernels.
 // Internally, kernels always operate on BHLD [batch, n_heads, seq_len, head_dim].
 // When the caller passes BLHD, dims 1 and 2 are transposed at entry.
 enum TensorLayout : int {
    BHLD = 0,  // [batch, n_heads, seq_len, head_dim]
    BLHD = 1,  // [batch, seq_len, n_heads, head_dim]
 };
 // Unified attention params covering BOTH addressing modes:
 //   - Contiguous K/V: dense [batch, kv_head, kv_len, head_dim] tensors (k/v).
 //   - Paged (SGLang-style): flat pool [size, kv_head, head_dim] + req_to_token.
 // Each kernel selects the addressing via a KVSource policy (see
 // attn_kv_source.cuh); a given call only touches the fields of one mode, so
 // this is a POD shared by both paths rather than two parallel structs that
 // drift out of sync.
 template<typename T, typename AT = float>
 struct AttentionParams {
    // ---- shared across all paths ----
    int batch;
    int q_head;
    int kv_head;
    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;
    // Mask: 2D [batch, kv_len], 3D [batch, q_len, kv_len],
    // or 4D [batch, n_heads, q_len, kv_len] (head dim broadcasts when stride=0)
    int mask_b_stride;   // batch stride
    int mask_h_stride;   // head stride (0 = broadcast across heads)
    int mask_q_stride;   // q stride (0 = all q rows share)
    const bool* __restrict__ mask;
    const T* __restrict__ q;
    T* __restrict__ o;
    AT* __restrict__ o_part;
    AT* __restrict__ ml_part;
    // ---- contiguous K/V mode ----
    int q_len;
    int kv_len;
    int kv_stride_b, kv_stride_h, kv_stride_l, kv_stride_d;
    const T* __restrict__ k;
    const T* __restrict__ v;
    // ---- paged (SGLang flat pool) mode ----
    const T* __restrict__ k_cache;
    const T* __restrict__ v_cache;
    // Indexing
    const int64_t* __restrict__ req_to_token;      // [num_reqs, max_context_len]
    const int64_t* __restrict__ req_pool_indices;   // [batch]
    const int* __restrict__ kv_indptr;               // [batch+1]
    const int* __restrict__ qo_indptr;               // [batch+1] or nullptr (decode)
    int max_context_len;  // req_to_token stride (dim 1)
    int max_seq_len;      // max per-request seq_len (host-side, for split computation)
    int total_q;          // total Q tokens across all requests (host-side, for grid)
    int max_q_len;        // max per-request q_len (host-side, for prefill grid)
 };
@@ -0,0 +1,56 @@
 #include "attn_dispatchers.cuh"
 #include "attn_entry_utils.cuh"
 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,
    c10::optional<torch::Tensor> o_part_buf,
    c10::optional<torch::Tensor> ml_part_buf
 ) {
    const at::cuda::OptionalCUDAGuard device_guard(device_of(q));
    auto stream = at::cuda::getCurrentCUDAStream();
    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");
    auto O = torch::empty_strided(q.sizes(), q.strides(), q.options());
    auto O_view = (layout == BLHD) ? O.transpose(1, 2) : O;
    p.o = (bf16*)O_view.data_ptr();
    if (o_part_buf.has_value() && ml_part_buf.has_value()
        && o_part_buf->defined() && ml_part_buf->defined()) {
        TORCH_CHECK(o_part_buf->scalar_type() == torch::kFloat32, "o_part_buf must be f32");
        TORCH_CHECK(ml_part_buf->scalar_type() == torch::kFloat32, "ml_part_buf must be f32");
        int64_t o_needed = (int64_t)p.batch * p.q_head * MAX_SPLITS * p.head_dim;
        TORCH_CHECK(o_part_buf->numel() >= o_needed,
                     "o_part_buf too small: need ", o_needed, " got ", o_part_buf->numel());
        p.o_part = (float*)o_part_buf->data_ptr();
        p.ml_part = (float*)ml_part_buf->data_ptr();
    } else {
        alloc_split_partials(p);
    }
    DISPATCH_HEAD_DIM(p.head_dim, dispatch_decode, p, stream);
    C10_CUDA_CHECK(cudaGetLastError());
    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") = (int64_t)BHLD,
        py::arg("o_part_buf") = py::none(),
        py::arg("ml_part_buf") = py::none(),
        "GQA decode (tensor-core head-packing on sm_80+, scalar fallback)");
 }
@@ -0,0 +1,143 @@
 #pragma once
 #include <cuda_bf16.h>
 #include <float.h>
 #include "attn_common.h"
 #include "attn_kv_source.cuh"
 #include "attn_warp_utils.cuh"
 constexpr int DC_CHUNK = 64;
 // Scalar split-KV decode (fallback for sm < 80, no tensor cores), unified
 // across contiguous and paged (SGLang flat-pool) K/V via the KV template
 // parameter.  For decode the query is the last token, so its valid range
 // [0, seq_len) IS the causal range; KV::decode_attend_len expresses that
 // bound per addressing mode (contig clips to causal_offset, paged = seq_len).
 template <int HEAD_DIM, typename KV, bool IsCausal, bool HasMask>
 __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;
    const int seq_len = KV::kv_len(p, batch);
    const KVContext kctx = KV::template make_ctx<HEAD_DIM>(p, batch, kv_head);
    // Q: [batch, q_head, q_len=1, head_dim] — stride-based
    float q_reg[8];
    int q_off = KV::q_decode_base(p, batch, q_head)
              + 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]);
    int mask_base = batch * p.mask_b_stride + q_head * p.mask_h_stride;
    float m = -FLT_MAX, d = 0.0f, acc_reg[8] = {0.0f};
    extern __shared__ __align__(16) bf16 smem[];
    bf16* k_smem = smem;
    bf16* v_smem = smem + DC_CHUNK * p.head_dim;
    // Split-KV: each split processes a contiguous subset of chunks
    int chunks_total = (seq_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, seq_len - chunk_start);
        // Load K and V into shared memory (addressing via KV policy;
        // paged guards empty slots with zero-fill).
        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 kc = chunk_start + s;
            KVAddr a = KV::kv_addr(p, kctx, kc, d_dim, true);
            k_smem[i] = a.valid ? *reinterpret_cast<const bf16*>(a.k) : (bf16)0.f;
            v_smem[i] = a.valid ? *reinterpret_cast<const bf16*>(a.v) : (bf16)0.f;
        }
        __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 constexpr (HasMask) {
                if (!p.mask[mask_base + kv_idx])
                    partial = -FLT_MAX;
            }
            if constexpr (IsCausal) {
                if (kv_idx >= KV::decode_attend_len(p, batch))
                    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;
            for (int i = 0; i < hd_per_thread; i++) {
                float vv = __bfloat162float(v_smem[s * p.head_dim + lane * hd_per_thread + i]);
                acc_reg[i] = fmaf(acc_reg[i], alpha, vv * 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 * MAX_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;
    }
 }
 // Split-combine: merges the per-split partials (o_part/ml_part) into the
 // final normalised O.  KV selects the O addressing (contig batch stride vs
 // paged row stride).
 template <typename KV>
 __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 * MAX_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 = fmaf(acc, corr, op[s * p.head_dim + d] * e);
        l = fmaf(l, corr, li * e);
        m = nm;
    }
    float inv = (l > 1e-20f) ? (1.0f / l) : 0.0f;
    int o_off = KV::q_decode_base(p, batch, q_head) + 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_kv_source.cuh"
 #include "attn_mma_utils.cuh"
 #include "attn_warp_utils.cuh"
 // Split-K (FlashDecoding) tensor-core decode via GQA head-packing, unified
 // across contiguous and paged (SGLang flat-pool) K/V via the KV template
 // parameter.  Decode has q_len == 1, so we pack G = q_head/kv_head query
 // 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.
 //
 // KV = ContigKV (dense tensors) or PagedKV (flat pool + req_to_token).
 // IsCausal and HasMask are compile-time bools — no runtime branch in the
 // inner compute loop.
 //
 // Traits = KernelTraits<HEAD_DIM, BC=16, WARPS=1, STAGES=2>.
 template <typename Traits, typename KV, bool IsCausal, bool HasMask>
 __global__ void attn_decode_split_kv_mma_kernel(AttentionParams<bf16> p) {
    const int lane = threadIdx.x;
    const int gid = lane >> 2;
    const int tid4 = lane & 3;
    const int pass = blockIdx.x / p.kv_head;
    const int kv_head = blockIdx.x % p.kv_head;
    const int batch = blockIdx.y;
    const int split = blockIdx.z;
    constexpr int MAX_G = 16;
    const int G_total = p.q_head / p.kv_head;
    const int g_begin = pass * MAX_G;
    const int G = min(MAX_G, G_total - g_begin);
    const int q_head0 = kv_head * G_total + g_begin;
    // Per-request seq_len (paged reads kv_indptr; contig uses p.kv_len).
    const int seq_len = KV::kv_len(p, batch);
    const KVContext kctx = KV::template make_ctx<Traits::HEAD_DIM>(p, batch, kv_head);
    // Double-buffered shared memory for K/V (no sQ needed)
    __shared__ __align__(16) bf16 sK[Traits::STAGES * Traits::BC * Traits::LD];
    __shared__ __align__(16) bf16 sV[Traits::STAGES * Traits::BC * Traits::LD];
    // Load Q directly from global into mma A-operand registers.
    const int q_base = KV::q_decode_base(p, batch, q_head0);
    const int qra = gid;
    const int qrb = gid + 8;
    const bool va = qra < G, vb = qrb < G;
    unsigned Qa[Traits::KD][4];
    load_q_mma_frags<Traits::KD>(p.q + q_base, p.q_stride_h, p.q_stride_d,
                                  qra, qrb, va, vb, tid4, Qa);
    float Oacc[Traits::DN8][4];
    #pragma unroll
    for (int j = 0; j < Traits::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 = (seq_len + Traits::BC - 1) / Traits::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);
    // ---- Load tile lambda: predicated cp.async (addressing via KV policy) ----
    auto load_tile = [&](int ti, int buf) {
        int kv0 = ti * Traits::BC;
        bf16* dK = sK + buf * Traits::BC * Traits::LD;
        bf16* dV = sV + buf * Traits::BC * Traits::LD;
        #pragma unroll
        for (int i = lane * Traits::VEC; i < Traits::TOTAL;
             i += Traits::NUM_THREADS * Traits::VEC) {
            int r = i / Traits::HEAD_DIM, d = i % Traits::HEAD_DIM;
            int kc = kv0 + r;
            bool valid = kc < seq_len;
            KVAddr a = KV::kv_addr(p, kctx, kc, d, valid);
            int off = r * Traits::LD + swiz_col(d, r, Traits::SWIZ_MASK);
            cp_async_16_pred(&dK[off], a.k, a.valid);
            cp_async_16_pred(&dV[off], a.v, a.valid);
        }
        cp_async_commit();
    };
    // ---- Multi-stage cp.async pipeline ----
    // Prologue loads STAGES tiles; each loop iteration waits only for the
    // oldest outstanding group (wait_group<STAGES-1>) so the STAGES-1 newer
    // tile loads stay in flight and overlap with the current tile's compute.
    constexpr int STAGES = Traits::STAGES;
    const int ntiles = ti_end - ti_begin;
    auto process_tile = [&](int it, int buf) {
        const bf16* bK = sK + buf * Traits::BC * Traits::LD;
        const bf16* bV = sV + buf * Traits::BC * Traits::LD;
        int kv0 = (ti_begin + it) * Traits::BC;
        float Sacc[Traits::NC8][4];
        mma_compute_scores<Traits>(Qa, bK, lane, Sacc);
        #pragma unroll
        for (int n8 = 0; n8 < Traits::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.  Paged treats [0, seq_len) as
        // the causal range (query is the last token); contig clips to the
        // causal_offset bound.  Dead code eliminated when IsCausal == false.
        int maxc = IsCausal ? KV::decode_attend_len(p, batch) : seq_len;
        mma_softmax_tile<Traits, HasMask>(kv0, maxc, maxc,
                                           0, 0,
                                           p.mask_b_stride, 0, 0,
                                           batch, 0,
                                           p.mask,
                                           Sacc, Oacc, m0, m1, l0, l1, lane);
        mma_pv_accumulate<Traits>(Sacc, bV, lane, Oacc);
    };
    if (ntiles >= STAGES) {
        #pragma unroll
        for (int i = 0; i < STAGES; i++)
            load_tile(ti_begin + i, i);
        for (int it = 0; it < ntiles; it++) {
            cp_async_wait_group<STAGES - 1>();
            __syncwarp();
            process_tile(it, it & (STAGES - 1));
            __syncwarp();
            if (it + STAGES < ntiles)
                load_tile(ti_begin + it + STAGES, (it + STAGES) & (STAGES - 1));
        }
    } else {
        // Fewer tiles than stages: load all, wait for all, process.
        for (int i = 0; i < ntiles; i++)
            load_tile(ti_begin + i, i);
        cp_async_wait_group<0>();
        __syncwarp();
        for (int it = 0; it < ntiles; it++)
            process_tile(it, it);
    }
    // ---- 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 * MAX_SPLITS + split;
    };
    #pragma unroll
    for (int dn8 = 0; dn8 < Traits::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) * Traits::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) * Traits::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,205 @@
 #pragma once
 // Shared attention dispatchers — used by both production .cu and test .cu.
 // No torch dependency; pure CUDA.
 //
 // The paged and contiguous kernels are unified by the KVSource policy
 // (ContigKV / PagedKV from attn_kv_source.cuh), so each launcher struct
 // below is templated on KV and the paged dispatch is just the same launcher
 // instantiated with PagedKV.  Only the grid/split math differs, and that is
 // covered by KV::host_q_len / KV::host_kv_len.
 #include <cuda_runtime.h>
 #include <algorithm>
 #include "attn_warp_utils.cuh"
 #include "attn_kv_source.cuh"
 #include "attn_prefill_split_q.cuh"
 #include "attn_decode_split_kv.cuh"
 #ifndef ASTRAI_NO_MMA
 #include "attn_prefill_split_q_mma.cuh"
 #include "attn_decode_split_kv_mma.cuh"
 #endif
 // Split-KV: compute number of splits to fill all SMs for small-batch decode.
 // Caps splits so each split processes at least `min_tiles_per_split` tiles,
 // avoiding excessive loop/prologue overhead when tiles are small.
 //
 // Target total grid blocks (`TARGET_BLOCKS`) rather than scaling splits by SM
 // count.  Decode blocks are single-warp (32 threads) and a SM hosts ~11 of
 // them, so the old `2*sm/base` cap badly undersplit at large batch (B=16 got
 // 3 splits, optimal ~8).  Measured (L20, grid search): bandwidth saturates
 // near 256-512 total blocks; 512 minimizes worst-case latency across the
 // B x kv grid; more is pure oversplit overhead.
 constexpr int DECODE_TARGET_BLOCKS = 512;
 inline int compute_num_splits(int base_blocks, int tiles_total,
                               int min_tiles_per_split = 1) {
    int n = (DECODE_TARGET_BLOCKS + base_blocks - 1) / base_blocks;
    int max_by_work = tiles_total / min_tiles_per_split;
    return std::max(1, std::min(n, std::min(max_by_work, MAX_SPLITS)));
 }
 // Dispatch IsCausal × HasMask — eliminates the duplicated 4-way if/else
 // ladder that appeared in each dispatch_* function.  FN must be a function
 // template <int HEAD_DIM, bool IsCausal, bool HasMask>; HEAD_DIM is forwarded
 // as the first template argument so callers only spell it once.
 //
 // Usage:  DISPATCH_CAUSAL_MASK(is_causal, has_mask, launcher<KV>::template launch, HEAD_DIM, p, stream);
 #define DISPATCH_CAUSAL_MASK(is_causal, has_mask, FN, HEAD_DIM, ...) \
    do { \
        if (is_causal) { \
            if (has_mask) FN<HEAD_DIM, true,  true>(__VA_ARGS__); \
            else          FN<HEAD_DIM, true,  false>(__VA_ARGS__); \
        } else { \
            if (has_mask) FN<HEAD_DIM, false, true>(__VA_ARGS__); \
            else          FN<HEAD_DIM, false, false>(__VA_ARGS__); \
        } \
    } while (0)
 // ======================================================================
 // Prefill launchers (KV selects ContigKV or PagedKV addressing)
 // ======================================================================
 #ifndef ASTRAI_NO_MMA
 template <typename KV>
 struct PrefillLauncherMMA {
    template <int HEAD_DIM, bool IsCausal, bool HasMask>
    static void launch(AttentionParams<bf16>& p, cudaStream_t stream) {
        constexpr int WARPS = 4;
        constexpr int BC = (HEAD_DIM <= 128) ? 32 : 16;
        using Traits = KernelTraits<HEAD_DIM, BC, WARPS, 2>;
        int q_len = KV::host_q_len(p);
        dim3 grid((q_len + Traits::BR * WARPS - 1) / (Traits::BR * WARPS),
                  p.q_head, p.batch);
        dim3 block(Traits::NUM_THREADS);
        attn_prefill_split_q_mma_kernel<Traits, KV, IsCausal, HasMask>
            <<<grid, block, 0, stream>>>(p);
    }
 };
 #endif
 template <typename KV>
 struct PrefillLauncherScalar {
    template <int HEAD_DIM, bool IsCausal, bool HasMask>
    static void launch(AttentionParams<bf16>& p, cudaStream_t stream) {
        constexpr int G = 8, ROWS = 32, P_BC = 32;
        int q_len = KV::host_q_len(p);
        dim3 grid((q_len + ROWS - 1) / ROWS, p.q_head, p.batch);
        dim3 block(G, ROWS);
        attn_prefill_split_q_kernel_t<HEAD_DIM, KV, G, ROWS, P_BC, IsCausal, HasMask>
            <<<grid, block, 0, stream>>>(p);
    }
 };
 template <int HEAD_DIM>
 static inline void dispatch_prefill(AttentionParams<bf16>& p, cudaStream_t stream) {
    bool is_causal = (p.causal_offset >= 0);
    bool has_mask = (p.use_mask && p.mask);
 #ifndef ASTRAI_NO_MMA
    DISPATCH_CAUSAL_MASK(is_causal, has_mask,
                         PrefillLauncherMMA<ContigKV>::template launch,
                         HEAD_DIM, p, stream);
 #else
    DISPATCH_CAUSAL_MASK(is_causal, has_mask,
                         PrefillLauncherScalar<ContigKV>::template launch,
                         HEAD_DIM, p, stream);
 #endif
 }
 template <int HEAD_DIM>
 static inline void dispatch_paged_prefill(AttentionParams<bf16>& p, cudaStream_t stream) {
    bool is_causal = (p.causal_offset >= 0);
    bool has_mask = (p.use_mask && p.mask);
 #ifndef ASTRAI_NO_MMA
    DISPATCH_CAUSAL_MASK(is_causal, has_mask,
                         PrefillLauncherMMA<PagedKV>::template launch,
                         HEAD_DIM, p, stream);
 #else
    DISPATCH_CAUSAL_MASK(is_causal, has_mask,
                         PrefillLauncherScalar<PagedKV>::template launch,
                         HEAD_DIM, p, stream);
 #endif
 }
 // ======================================================================
 // Decode launchers (KV selects ContigKV or PagedKV addressing)
 // ======================================================================
 #ifndef ASTRAI_NO_MMA
 // BC=16: halves smem (16KB vs 32KB) → doubles occupancy (6 vs 3 blocks/SM).
 // For D=256, BC=16 also reduces register pressure (fewer Sacc/PV frags),
 // enabling STAGES=2 (double-buffer) within the 32KB smem budget — eliminates
 // the 176-byte spill that STAGES=1+BC=32 suffered.
 template <typename KV>
 struct DecodeLauncherMMA {
    template <int HEAD_DIM, bool IsCausal, bool HasMask>
    static void launch(AttentionParams<bf16>& p, int group_size, cudaStream_t stream) {
        int G = p.q_head / p.kv_head;
        constexpr int MAX_G = 16;
        int num_passes = (G + MAX_G - 1) / MAX_G;
        constexpr int BC = 16;
        int kv_len = KV::host_kv_len(p);
        int tiles_total = (kv_len + BC - 1) / BC;
        p.num_splits = compute_num_splits(p.batch * p.kv_head * num_passes, tiles_total, 2);
        constexpr int STAGES = 2;
        using Traits = KernelTraits<HEAD_DIM, BC, 1, STAGES>;
        dim3 grid(p.kv_head * num_passes, p.batch, p.num_splits);
        attn_decode_split_kv_mma_kernel<Traits, KV, IsCausal, HasMask>
            <<<grid, 32, 0, stream>>>(p);
    }
 };
 #endif
 template <typename KV>
 struct DecodeLauncherScalar {
    template <int HEAD_DIM, bool IsCausal, bool HasMask>
    static void launch(AttentionParams<bf16>& p, int group_size, cudaStream_t stream) {
        int kv_len = KV::host_kv_len(p);
        int chunks_total = (kv_len + DC_CHUNK - 1) / DC_CHUNK;
        p.num_splits = compute_num_splits(p.batch * p.kv_head, chunks_total);
        size_t smem = 2 * DC_CHUNK * p.head_dim * sizeof(bf16);
        int g = min(group_size, 32);  // cap at 32 to respect 1024-thread limit
        dim3 grid(p.batch * p.kv_head, 1, p.num_splits);
        dim3 block(32, g);
        attn_decode_split_kv_kernel<HEAD_DIM, KV, IsCausal, HasMask>
            <<<grid, block, smem, stream>>>(p);
    }
 };
 template <int HEAD_DIM>
 static inline void dispatch_decode(AttentionParams<bf16>& p, cudaStream_t stream) {
    bool is_causal = (p.causal_offset >= 0);
    bool has_mask = (p.use_mask && p.mask);
    int group_size = p.q_head / p.kv_head;
 #ifndef ASTRAI_NO_MMA
    DISPATCH_CAUSAL_MASK(is_causal, has_mask,
                         DecodeLauncherMMA<ContigKV>::template launch,
                         HEAD_DIM, p, group_size, stream);
 #else
    DISPATCH_CAUSAL_MASK(is_causal, has_mask,
                         DecodeLauncherScalar<ContigKV>::template launch,
                         HEAD_DIM, p, group_size, stream);
 #endif
    attn_decode_combine_kernel<ContigKV><<<p.batch * p.q_head, p.head_dim, 0, stream>>>(p);
 }
 template <int HEAD_DIM>
 static inline void dispatch_paged_decode(AttentionParams<bf16>& p, cudaStream_t stream) {
    bool is_causal = (p.causal_offset >= 0);
    bool has_mask = (p.use_mask && p.mask);
    int group_size = p.q_head / p.kv_head;
 #ifndef ASTRAI_NO_MMA
    DISPATCH_CAUSAL_MASK(is_causal, has_mask,
                         DecodeLauncherMMA<PagedKV>::template launch,
                         HEAD_DIM, p, group_size, stream);
 #else
    DISPATCH_CAUSAL_MASK(is_causal, has_mask,
                         DecodeLauncherScalar<PagedKV>::template launch,
                         HEAD_DIM, p, group_size, stream);
 #endif
    attn_decode_combine_kernel<PagedKV><<<p.batch * p.q_head, p.head_dim, 0, stream>>>(p);
 }
@@ -0,0 +1,311 @@
 #pragma once
 #include <float.h>
 #include <torch/extension.h>
 #include <c10/cuda/CUDAGuard.h>
 #include "attn_common.h"
 #include "attn_warp_utils.cuh"
 using bf16 = __nv_bfloat16;
 // Dispatch head_dim: shared macro — avoids C++20 lambda template syntax.
 // Usage: DISPATCH_HEAD_DIM(hd, fn, args...)
 //   Expands to: fn<32>(args...); fn<64>(args...); etc.
 #define DISPATCH_HEAD_DIM(hd, fn, ...) \
    switch (hd) { \
        case 32:  fn<32>(__VA_ARGS__); break; \
        case 64:  fn<64>(__VA_ARGS__); break; \
        case 128: fn<128>(__VA_ARGS__); break; \
        case 256: fn<256>(__VA_ARGS__); break; \
        default: \
            TORCH_CHECK(false, "unsupported head_dim ", hd, \
                         " (supported: 32, 64, 128, 256)"); \
    }
 // The split kernel unconditionally writes every (batch, q_head, split) slot it
 // owns — including empty split ranges, which store m = -FLT_MAX so the combine
 // skips them. Allocators are therefore left uninitialized (torch::empty); the
 // per-call memset (torch::zeros / torch::full) was pure overhead.
 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(at::IntArrayRef{p.batch, p.q_head, MAX_SPLITS, p.head_dim}, fopt);
    auto ml_part = torch::empty(at::IntArrayRef{p.batch, p.q_head, MAX_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 == BLHD) 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 ----
 // Accepts 2D [batch, kv_len], 3D [batch, q_len, kv_len],
 // or 4D [batch, n_heads, q_len, kv_len].
 // Head/q dimensions with size 1 broadcast (stride set to 0).
 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_h_stride = 0;
            p.mask_q_stride = 0;
        } else if (m.dim() == 3) {
            TORCH_CHECK(m.size(1) == 1 || m.size(1) == p.q_len, "mask q_len mismatch");
            p.mask_b_stride = (int)m.stride(0);
            p.mask_h_stride = 0;
            p.mask_q_stride = (m.size(1) == 1) ? 0 : (int)m.stride(1);
        } else if (m.dim() == 4) {
            TORCH_CHECK(m.size(2) == 1 || m.size(2) == p.q_len, "mask q_len mismatch");
            p.mask_b_stride = (int)m.stride(0);
            p.mask_h_stride = (m.size(1) == 1) ? 0 : (int)m.stride(1);
            p.mask_q_stride = (m.size(2) == 1) ? 0 : (int)m.stride(2);
        } else {
            TORCH_CHECK(false, "mask must be 2D, 3D, or 4D");
        }
        p.mask = m.data_ptr<bool>();
    } else {
        p.mask = nullptr;
        p.mask_b_stride = 0;
        p.mask_h_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 == BLHD) 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_decode_params ----
 // SGLang-style: flat KV pool + req_to_token indexing + variable
 // seq_lens via kv_indptr.  Q is [batch, q_head, head_dim] (q_len=1 per req).
 template<typename T>
 inline void attn_pack_paged_decode_params(
    torch::Tensor q,
    torch::Tensor k_cache,
    torch::Tensor v_cache,
    torch::Tensor req_to_token,
    torch::Tensor req_pool_indices,
    torch::Tensor kv_indptr,
    int64_t max_seq_len,
    c10::optional<torch::Tensor> mask,
    int64_t causal_offset,
    double scale,
    AttentionParams<T>& p
 ) {
    const at::cuda::OptionalCUDAGuard device_guard(device_of(q));
    TORCH_CHECK(q.is_cuda() && k_cache.is_cuda() && v_cache.is_cuda());
    TORCH_CHECK(req_to_token.is_cuda() && req_pool_indices.is_cuda() && kv_indptr.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(req_to_token.dtype() == torch::kLong, "req_to_token must be int64");
    TORCH_CHECK(req_pool_indices.dtype() == torch::kLong, "req_pool_indices must be int64");
    TORCH_CHECK(kv_indptr.dtype() == torch::kInt32, "kv_indptr must be int32");
    TORCH_CHECK(k_cache.sizes() == v_cache.sizes(), "k_cache and v_cache must match");
    TORCH_CHECK(k_cache.dim() == 3, "k_cache must be 3D [size, kv_head, head_dim]");
    TORCH_CHECK(q.dim() == 3, "q must be 3D [batch, q_head, head_dim]");
    p.batch = (int)q.size(0);
    p.q_head = (int)q.size(1);
    p.head_dim = (int)q.size(2);
    p.kv_head = (int)k_cache.size(1);
    TORCH_CHECK(k_cache.size(2) == p.head_dim, "k_cache head_dim mismatch");
    TORCH_CHECK(p.head_dim % 32 == 0, "head_dim must be multiple of 32");
    TORCH_CHECK(p.q_head % p.kv_head == 0, "q_head must be divisible by kv_head");
    p.q_stride_l = (int)q.stride(0);
    p.q_stride_h = (int)q.stride(1);
    p.q_stride_d = (int)q.stride(2);
    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.req_to_token = req_to_token.data_ptr<int64_t>();
    p.req_pool_indices = req_pool_indices.data_ptr<int64_t>();
    p.kv_indptr = kv_indptr.data_ptr<int>();
    p.qo_indptr = nullptr;
    p.max_context_len = (int)req_to_token.size(1);
    p.max_seq_len = (int)max_seq_len;
    p.total_q = p.batch;  // decode: 1 Q token per request
    p.max_q_len = 1;
    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);
    if (p.use_mask) {
        auto m = mask.value();
        TORCH_CHECK(m.is_cuda() && m.dtype() == torch::kBool, "mask must be bool CUDA");
        TORCH_CHECK(m.size(0) == p.batch, "mask batch mismatch");
        p.mask_b_stride = (int)m.stride(0);
        p.mask_h_stride = 0;
        p.mask_q_stride = 0;
        p.mask = m.data_ptr<bool>();
    } else {
        p.mask = nullptr;
        p.mask_b_stride = 0;
        p.mask_h_stride = 0;
        p.mask_q_stride = 0;
    }
    p.o = nullptr;
    p.o_part = nullptr;
    p.ml_part = nullptr;
 }
 // ---- attn_pack_paged_prefill_params ----
 // SGLang-style: flat KV pool + req_to_token + ragged batch via qo_indptr.
 // Q is [total_q, q_head, head_dim] (flattened across all requests).
 template<typename T>
 inline void attn_pack_paged_prefill_params(
    torch::Tensor q,
    torch::Tensor k_cache,
    torch::Tensor v_cache,
    torch::Tensor req_to_token,
    torch::Tensor req_pool_indices,
    torch::Tensor kv_indptr,
    torch::Tensor qo_indptr,
    c10::optional<torch::Tensor> mask,
    int64_t max_q_len,
    int64_t causal_offset,
    double scale,
    AttentionParams<T>& p
 ) {
    const at::cuda::OptionalCUDAGuard device_guard(device_of(q));
    TORCH_CHECK(q.is_cuda() && k_cache.is_cuda() && v_cache.is_cuda());
    TORCH_CHECK(req_to_token.is_cuda() && req_pool_indices.is_cuda());
    TORCH_CHECK(kv_indptr.is_cuda() && qo_indptr.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(req_to_token.dtype() == torch::kLong, "req_to_token must be int64");
    TORCH_CHECK(req_pool_indices.dtype() == torch::kLong, "req_pool_indices must be int64");
    TORCH_CHECK(kv_indptr.dtype() == torch::kInt32, "kv_indptr must be int32");
    TORCH_CHECK(qo_indptr.dtype() == torch::kInt32, "qo_indptr must be int32");
    TORCH_CHECK(k_cache.sizes() == v_cache.sizes(), "k_cache and v_cache must match");
    TORCH_CHECK(k_cache.dim() == 3, "k_cache must be 3D [size, kv_head, head_dim]");
    TORCH_CHECK(q.dim() == 3, "q must be 3D [total_q, q_head, head_dim]");
    p.q_head = (int)q.size(1);
    p.head_dim = (int)q.size(2);
    p.kv_head = (int)k_cache.size(1);
    p.batch = (int)req_pool_indices.size(0);
    TORCH_CHECK(k_cache.size(2) == p.head_dim, "k_cache head_dim mismatch");
    TORCH_CHECK(p.head_dim % 16 == 0, "head_dim must be multiple of 16");
    TORCH_CHECK(p.q_head % p.kv_head == 0, "q_head must be divisible by kv_head");
    TORCH_CHECK(kv_indptr.size(0) == p.batch + 1, "kv_indptr must be [batch+1]");
    TORCH_CHECK(qo_indptr.size(0) == p.batch + 1, "qo_indptr must be [batch+1]");
    p.q_stride_l = (int)q.stride(0);
    p.q_stride_h = (int)q.stride(1);
    p.q_stride_d = (int)q.stride(2);
    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.req_to_token = req_to_token.data_ptr<int64_t>();
    p.req_pool_indices = req_pool_indices.data_ptr<int64_t>();
    p.kv_indptr = kv_indptr.data_ptr<int>();
    p.qo_indptr = qo_indptr.data_ptr<int>();
    p.max_context_len = (int)req_to_token.size(1);
    p.total_q = (int)q.size(0);  // prefill: flattened Q across all requests
    p.max_q_len = (int)max_q_len;
    // max_seq_len is unused by the prefill path (decode uses it for split
    // computation); fill with max_q_len only to keep the POD struct defined.
    p.max_seq_len = p.max_q_len;
    p.causal_offset = (int)causal_offset;
    p.use_mask = (mask.has_value() && mask.value().defined()) ? 1 : 0;
    if (p.use_mask) {
        auto m = mask.value();
        TORCH_CHECK(m.is_cuda() && m.dtype() == torch::kBool, "mask must be bool CUDA");
        TORCH_CHECK(m.size(0) == p.batch, "mask batch mismatch");
        if (m.dim() == 2) {
            TORCH_CHECK(m.size(1) <= p.max_context_len, "mask kv_len mismatch");
            p.mask_b_stride = (int)m.stride(0);
            p.mask_h_stride = 0;
            p.mask_q_stride = 0;
        } else if (m.dim() == 4) {
            TORCH_CHECK(m.size(1) == 1 || m.size(1) == p.q_head, "mask head mismatch");
            TORCH_CHECK(m.size(2) == 1 || m.size(2) == p.max_q_len, "mask q_len mismatch");
            TORCH_CHECK(m.size(3) <= p.max_context_len, "mask kv_len mismatch");
            p.mask_b_stride = (int)m.stride(0);
            p.mask_h_stride = (m.size(1) == 1) ? 0 : (int)m.stride(1);
            p.mask_q_stride = (m.size(2) == 1) ? 0 : (int)m.stride(2);
        } else {
            TORCH_CHECK(false, "mask must be 2D or 4D");
        }
        p.mask = m.data_ptr<bool>();
    } else {
        p.mask = nullptr;
        p.mask_b_stride = 0;
        p.mask_h_stride = 0;
        p.mask_q_stride = 0;
    }
    p.scale = (scale > 0.0) ? (float)scale : 1.0f / sqrtf((float)p.head_dim);
    p.o = nullptr;
    p.o_part = nullptr;
    p.ml_part = nullptr;
 }
@@ -0,0 +1,159 @@
 #pragma once
 #include <cuda_bf16.h>
 #include "attn_common.h"
 // ============================================================================
 // KVSource policies — the single dimension along which the paged and
 // non-paged attention kernels differ.  Each kernel is templated on one of
 // these (ContigKV / PagedKV) and stays fully generic: the policy owns every
 // place where "where does K/V live" and "what is this request's seq_len"
 // are answered.  All methods are __host__ __device__ so the same policy
 // serves both the device kernels (addressing, seq_len) and the host-side
 // launchers (grid / split computation).
 //
 //   ContigKV:  K/V are dense [batch, kv_head, kv_len, head_dim] tensors.
 //              Params fields used: k, v, kv_stride_*, kv_len, q_len,
 //              q_stride_b, causal_offset.
 //   PagedKV:   K/V live in a flat pool [size, kv_head, head_dim] indexed via
 //              req_to_token.  Params fields used: k_cache, v_cache,
 //              req_to_token, req_pool_indices, kv_indptr, qo_indptr,
 //              max_context_len, q_stride_l.
 //
 // Addressing state that is constant across a whole kernel invocation for one
 // (batch, kv_head) pair is captured once by make_ctx<HEAD_DIM>() and passed
 // to kv_addr, so the load loops never redo the hoistable base computation
 // (e.g. the req_pool_indices global read) element-by-element.
 // ============================================================================
 // Every policy method is static + callable from both host and device code.
 #define HOST_DEV_FORCEINLINE static __host__ __device__ __forceinline__
 using bf16 = __nv_bfloat16;
 // Hoisted per-(batch, kv_head) addressing context.
 struct KVContext {
    int kv_base;          // contig: batch*kv_stride_b + kv_head*kv_stride_h
    int64_t req_idx;      // paged: req_pool_indices[batch]
    int64_t rtt_stride;   // paged: max_context_len
    int64_t pool_stride;  // paged: kv_head * HEAD_DIM
    int64_t head_off;     // paged: kv_head * HEAD_DIM
 };
 // Per-element K/V global addresses for one (kc, d) position of a K/V tile.
 // The pointers are ALWAYS the computed addresses (never nullptr) — callers
 // gate on `valid` (cp.async src_size=0, or a guarded scalar deref).  `valid`
 // starts as "within the request's seq_len"; the paged policy further degrades
 // it when req_to_token maps the position to a negative slot (empty padding).
 // This matches the original hand-rolled load loops, where the address was
 // always formed and the predicate decided whether anything was read.
 struct KVAddr {
    const void* k;
    const void* v;
    bool valid;
 };
 // ---- Contiguous K/V ----
 struct ContigKV {
    static constexpr bool kPaged = false;
    // host-side length hooks (grid + split computation in the launchers)
    HOST_DEV_FORCEINLINE int host_q_len(const AttentionParams<bf16>& p) {
        return p.q_len;
    }
    HOST_DEV_FORCEINLINE int host_kv_len(const AttentionParams<bf16>& p) {
        return p.kv_len;
    }
    // prefill: element offset of the request's Q rows (kernel adds qrow*q_stride_l)
    HOST_DEV_FORCEINLINE int q_base(
        const AttentionParams<bf16>& p, int batch, int q_head) {
        return batch * p.q_stride_b + q_head * p.q_stride_h;
    }
    // decode: same offset (q_len == 1, so there is no row stride component)
    HOST_DEV_FORCEINLINE int q_decode_base(
        const AttentionParams<bf16>& p, int batch, int q_head) {
        return batch * p.q_stride_b + q_head * p.q_stride_h;
    }
    HOST_DEV_FORCEINLINE int kv_len(const AttentionParams<bf16>& p, int batch) {
        return p.kv_len;
    }
    HOST_DEV_FORCEINLINE int q_len(const AttentionParams<bf16>& p, int batch) {
        return p.q_len;
    }
    HOST_DEV_FORCEINLINE int causal_offset(const AttentionParams<bf16>& p, int batch) {
        return p.causal_offset;
    }
    // decode: exclusive bound of the single query's attend range
    HOST_DEV_FORCEINLINE int decode_attend_len(const AttentionParams<bf16>& p, int batch) {
        return (p.kv_len < p.causal_offset + 1) ? p.kv_len : (p.causal_offset + 1);
    }
    template <int HEAD_DIM>
    HOST_DEV_FORCEINLINE KVContext make_ctx(
        const AttentionParams<bf16>& p, int batch, int kv_head) {
        KVContext c = {};
        c.kv_base = batch * p.kv_stride_b + kv_head * p.kv_stride_h;
        return c;
    }
    HOST_DEV_FORCEINLINE KVAddr kv_addr(
        const AttentionParams<bf16>& p, const KVContext& c, int kc, int d, bool valid) {
        const int g_off = c.kv_base + kc * p.kv_stride_l + d * p.kv_stride_d;
        return {&p.k[g_off], &p.v[g_off], valid};
    }
 };
 // ---- Paged (SGLang-style flat pool) K/V ----
 struct PagedKV {
    static constexpr bool kPaged = true;
    HOST_DEV_FORCEINLINE int host_q_len(const AttentionParams<bf16>& p) {
        return p.max_q_len;
    }
    HOST_DEV_FORCEINLINE int host_kv_len(const AttentionParams<bf16>& p) {
        return p.max_context_len;
    }
    // prefill: Q rows start at qo_indptr[batch] (ragged batch base)
    HOST_DEV_FORCEINLINE int q_base(
        const AttentionParams<bf16>& p, int batch, int q_head) {
        return p.qo_indptr[batch] * p.q_stride_l + q_head * p.q_stride_h;
    }
    // decode: Q is [batch, q_head, head_dim], so batch is the outer row
    HOST_DEV_FORCEINLINE int q_decode_base(
        const AttentionParams<bf16>& p, int batch, int q_head) {
        return batch * p.q_stride_l + q_head * p.q_stride_h;
    }
    HOST_DEV_FORCEINLINE int kv_len(const AttentionParams<bf16>& p, int batch) {
        return p.kv_indptr[batch + 1] - p.kv_indptr[batch];
    }
    HOST_DEV_FORCEINLINE int q_len(const AttentionParams<bf16>& p, int batch) {
        return p.qo_indptr[batch + 1] - p.qo_indptr[batch];
    }
    HOST_DEV_FORCEINLINE int causal_offset(const AttentionParams<bf16>& p, int batch) {
        return kv_len(p, batch) - q_len(p, batch);
    }
    // decode: the query is the last token, so [0, seq_len) IS its causal range
    HOST_DEV_FORCEINLINE int decode_attend_len(const AttentionParams<bf16>& p, int batch) {
        return kv_len(p, batch);
    }
    template <int HEAD_DIM>
    HOST_DEV_FORCEINLINE KVContext make_ctx(
        const AttentionParams<bf16>& p, int batch, int kv_head) {
        KVContext c = {};
        c.req_idx = p.req_pool_indices[batch];
        c.rtt_stride = (int64_t)p.max_context_len;
        c.pool_stride = (int64_t)p.kv_head * HEAD_DIM;
        c.head_off = (int64_t)kv_head * HEAD_DIM;
        return c;
    }
    HOST_DEV_FORCEINLINE KVAddr kv_addr(
        const AttentionParams<bf16>& p, const KVContext& c, int kc, int d, bool valid) {
        const int64_t slot = valid ? p.req_to_token[c.req_idx * c.rtt_stride + kc] : 0;
        const bool ok = valid && (slot >= 0);
        const int64_t gmem_off = slot * c.pool_stride + c.head_off + d;
        return {&p.k_cache[gmem_off], &p.v_cache[gmem_off], ok};
    }
 };
--- a/Show More
+++ b/Show More
		`@@ -0,0 +1 @@`
							"""Compiled CUDA kernel modules (``*.so``) live here, kept separate from Python source."""
		`@@ -0,0 +1,2 @@`
							`# Source directory for CUDA kernels — build-time only.`
							`# Compiled .so files live in astrAI/_ext/.`