refactor: standardize packed 3d inference

- keep training attention on dense 4d tensors
- use packed 3d tensors with KV cache for inference
- extend CUDA rotary embedding to packed 3d inputs
- adapt torch, CUDA and FlashAttention backend dispatch

Benchmark: NVIDIA L20, BF16, 1B model, paged KV cache, CUDA Graph, prompt 512, generation 128 (median of 3 alternating runs)
- batch 1: 234.5 -> 242.6 tok/s (1.034x, +3.4%)
- batch 8: 1243.1 -> 1286.6 tok/s (1.035x, +3.5%)

.github/workflows/release.yml

@@ -54,6 +54,9 @@ jobs:
       - name: Build wheel (with CUDA kernels)
         run: |
           CSRC_KERNELS=true pip wheel . --no-deps --no-build-isolation -w dist/
+          for f in dist/*.whl; do
+            mv "$f" "dist/$(basename "$f" .whl)+${{ matrix.cuda_tag }}.whl"
+          done
 
       - uses: actions/upload-artifact@v4
         with:

.gitignore

@@ -9,6 +9,7 @@
 !scripts/**/*.py
 !tests/**/*.py
 !csrc/**/*.py
+!csrc/CMakeLists.txt
 
 !csrc/**/*.cu
 !csrc/**/*.h

CONTRIBUTING.md

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

Dockerfile

@@ -57,8 +57,13 @@ COPY docs/ ./docs/
 COPY pyproject.toml .
 COPY README.md .
 
-# Create non-root user
-RUN useradd -m astrai && chown -R astrai:astrai /app
+# Create non-root user matching the host uid/gid (passed via build args)
+ARG USER_UID=1000
+ARG USER_GID=1000
+RUN groupadd -g "${USER_GID}" astrai \
+    && useradd -m -u "${USER_UID}" -g astrai astrai \
+    && chown -R astrai:astrai /app
+ENV HOME=/home/astrai
 USER astrai
 
 ENV PYTHONUNBUFFERED=1 \

LICENSE

@@ -1,674 +1,201 @@
-                    GNU GENERAL PUBLIC LICENSE
-                       Version 3, 29 June 2007
-
- Copyright (C) 2007 Free Software Foundation, Inc. <https://fsf.org/>
- Everyone is permitted to copy and distribute verbatim copies
- of this license document, but changing it is not allowed.
-
-                            Preamble
-
-  The GNU General Public License is a free, copyleft license for
-software and other kinds of works.
-
-  The licenses for most software and other practical works are designed
-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
-share and change all versions of a program--to make sure it remains free
-software for all its users.  We, the Free Software Foundation, use the
-GNU General Public License for most of our software; it applies also to
-any other work released this way by its authors.  You can apply it to
-your programs, too.
-
-  When we speak of free software, we are referring to freedom, not
-price.  Our General Public Licenses are designed to make sure that you
-have the freedom to distribute copies of free software (and charge for
-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
-free programs, and that you know you can do these things.
-
-  To protect your rights, we need to prevent others from denying you
-these rights or asking you to surrender the rights.  Therefore, you have
-certain responsibilities if you distribute copies of the software, or if
-you modify it: responsibilities to respect the freedom of others.
-
-  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
-freedoms that you received.  You must make sure that they, too, receive
-or can get the source code.  And you must show them these terms so they
-know their rights.
-
-  Developers that use the GNU GPL protect your rights with two steps:
-(1) assert copyright on the software, and (2) offer you this License
-giving you legal permission to copy, distribute and/or modify it.
-
-  For the developers' and authors' protection, the GPL clearly explains
-that there is no warranty for this free software.  For both users' and
-authors' sake, the GPL requires that modified versions be marked as
-changed, so that their problems will not be attributed erroneously to
-authors of previous versions.
-
-  Some devices are designed to deny users access to install or run
-modified versions of the software inside them, although the manufacturer
-can do so.  This is fundamentally incompatible with the aim of
-protecting users' freedom to change the software.  The systematic
-pattern of such abuse occurs in the area of products for individuals to
-use, which is precisely where it is most unacceptable.  Therefore, we
-have designed this version of the GPL to prohibit the practice for those
-products.  If such problems arise substantially in other domains, we
-stand ready to extend this provision to those domains in future versions
-of the GPL, as needed to protect the freedom of users.
-
-  Finally, every program is threatened constantly by software patents.
-States should not allow patents to restrict development and use of
-software on general-purpose computers, but in those that do, we wish to
-avoid the special danger that patents applied to a free program could
-make it effectively proprietary.  To prevent this, the GPL assures that
-patents cannot be used to render the program non-free.
-
-  The precise terms and conditions for copying, distribution and
-modification follow.
-
-                       TERMS AND CONDITIONS
-
-  0. Definitions.
-
-  "This License" refers to version 3 of the GNU General Public License.
-
-  "Copyright" also means copyright-like laws that apply to other kinds of
-works, such as semiconductor masks.
-
-  "The Program" refers to any copyrightable work licensed under this
-License.  Each licensee is addressed as "you".  "Licensees" and
-"recipients" may be individuals or organizations.
-
-  To "modify" a work means to copy from or adapt all or part of the work
-in a fashion requiring copyright permission, other than the making of an
-exact copy.  The resulting work is called a "modified version" of the
-earlier work or a work "based on" the earlier work.
-
-  A "covered work" means either the unmodified Program or a work based
-on the Program.
-
-  To "propagate" a work means to do anything with it that, without
-permission, would make you directly or secondarily liable for
-infringement under applicable copyright law, except executing it on a
-computer or modifying a private copy.  Propagation includes copying,
-distribution (with or without modification), making available to the
-public, and in some countries other activities as well.
-
-  To "convey" a work means any kind of propagation that enables other
-parties to make or receive copies.  Mere interaction with a user through
-a computer network, with no transfer of a copy, is not conveying.
-
-  An interactive user interface displays "Appropriate Legal Notices"
-to the extent that it includes a convenient and prominently visible
-feature that (1) displays an appropriate copyright notice, and (2)
-tells the user that there is no warranty for the work (except to the
-extent that warranties are provided), that licensees may convey the
-work under this License, and how to view a copy of this License.  If
-the interface presents a list of user commands or options, such as a
-menu, a prominent item in the list meets this criterion.
-
-  1. Source Code.
-
-  The "source code" for a work means the preferred form of the work
-for making modifications to it.  "Object code" means any non-source
-form of a work.
-
-  A "Standard Interface" means an interface that either is an official
-standard defined by a recognized standards body, or, in the case of
-interfaces specified for a particular programming language, one that
-is widely used among developers working in that language.
-
-  The "System Libraries" of an executable work include anything, other
-than the work as a whole, that (a) is included in the normal form of
-packaging a Major Component, but which is not part of that Major
-Component, and (b) serves only to enable use of the work with that
-Major Component, or to implement a Standard Interface for which an
-implementation is available to the public in source code form.  A
-"Major Component", in this context, means a major essential component
-(kernel, window system, and so on) of the specific operating system
-(if any) on which the executable work runs, or a compiler used to
-produce the work, or an object code interpreter used to run it.
-
-  The "Corresponding Source" for a work in object code form means all
-the source code needed to generate, install, and (for an executable
-work) run the object code and to modify the work, including scripts to
-control those activities.  However, it does not include the work's
-System Libraries, or general-purpose tools or generally available free
-programs which are used unmodified in performing those activities but
-which are not part of the work.  For example, Corresponding Source
-includes interface definition files associated with source files for
-the work, and the source code for shared libraries and dynamically
-linked subprograms that the work is specifically designed to require,
-such as by intimate data communication or control flow between those
-subprograms and other parts of the work.
-
-  The Corresponding Source need not include anything that users
-can regenerate automatically from other parts of the Corresponding
-Source.
-
-  The Corresponding Source for a work in source code form is that
-same work.
-
-  2. Basic Permissions.
-
-  All rights granted under this License are granted for the term of
-copyright on the Program, and are irrevocable provided the stated
-conditions are met.  This License explicitly affirms your unlimited
-permission to run the unmodified Program.  The output from running a
-covered work is covered by this License only if the output, given its
-content, constitutes a covered work.  This License acknowledges your
-rights of fair use or other equivalent, as provided by copyright law.
-
-  You may make, run and propagate covered works that you do not
-convey, without conditions so long as your license otherwise remains
-in force.  You may convey covered works to others for the sole purpose
-of having them make modifications exclusively for you, or provide you
-with facilities for running those works, provided that you comply with
-the terms of this License in conveying all material for which you do
-not control copyright.  Those thus making or running the covered works
-for you must do so exclusively on your behalf, under your direction
-and control, on terms that prohibit them from making any copies of
-your copyrighted material outside their relationship with you.
-
-  Conveying under any other circumstances is permitted solely under
-the conditions stated below.  Sublicensing is not allowed; section 10
-makes it unnecessary.
-
-  3. Protecting Users' Legal Rights From Anti-Circumvention Law.
-
-  No covered work shall be deemed part of an effective technological
-measure under any applicable law fulfilling obligations under article
-11 of the WIPO copyright treaty adopted on 20 December 1996, or
-similar laws prohibiting or restricting circumvention of such
-measures.
-
-  When you convey a covered work, you waive any legal power to forbid
-circumvention of technological measures to the extent such circumvention
-is effected by exercising rights under this License with respect to
-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
-users, your or third parties' legal rights to forbid circumvention of
-technological measures.
-
-  4. Conveying Verbatim Copies.
-
-  You may convey verbatim copies of the Program's source code as you
-receive it, in any medium, provided that you conspicuously and
-appropriately publish on each copy an appropriate copyright notice;
-keep intact all notices stating that this License and any
-non-permissive terms added in accord with section 7 apply to the code;
-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>.
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

README.md

@@ -8,7 +8,7 @@
 
 <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/tag/ViperEkura/AstrAI?label=Release&color=76bad9" alt="release">
   <img src="https://img.shields.io/github/stars/ViperEkura/AstrAI?style=flat&label=Stars&color=76bad9" alt="stars">
   <img src="https://img.shields.io/github/forks/ViperEkura/AstrAI?style=flat&label=Forks&color=76bad9" alt="forks">
@@ -27,7 +27,7 @@
 
 ## 📖 Table of Contents
 
-- [Features](#features)
+- [Overview](#overview)
 - [Getting Started](#getting-started)
 - [Demo](#demo)
 - [Documentation](#documentation)
@@ -40,15 +40,19 @@
 <a id="english"></a>
 ## English
 
-### Features
+### Overview
 
-- 🚀 **High Performance**: Optimized for both training and inference with efficient parallelization.
-- 🔧 **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.
+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.
+
+| 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 |
 
 ### Getting Started
 
@@ -56,6 +60,8 @@ 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
@@ -132,7 +138,7 @@ Check out the demos in the `scripts/demo/` folder:
 # Download model weights (required before running demos)
 python scripts/demo/download.py                      # model → params/
 
-# Interactive streaming chat (multi-turn, maintains history)
+# Single-turn interactive streaming prompt loop (no conversation history)
 python scripts/demo/stream_chat.py
 # Type your message after >>, type !exit to quit
 
@@ -183,7 +189,7 @@ 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
 ```
 
@@ -250,10 +256,10 @@ For major changes, please open an issue first to discuss what you would like to
 
 ### 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>

astrai/__init__.py

@@ -1,9 +1,6 @@
-__version__ = "1.3.12"
+__version__ = "1.3.13"
 __author__ = "ViperEkura"
 
-import logging
-import os
-
 from astrai.config import (
     AutoRegressiveLMConfig,
     BaseModelConfig,
@@ -21,7 +18,6 @@ from astrai.dataset import (
 )
 from astrai.factory import BaseFactory
 from astrai.inference import (
-    GenerationRequest,
     InferenceEngine,
     ProtocolHandler,
     SamplingPipeline,
@@ -29,6 +25,7 @@ from astrai.inference import (
     run_server,
     sample,
 )
+from astrai.logging import setup_logging
 from astrai.model import (
     AutoModel,
     AutoRegressiveLM,
@@ -56,30 +53,6 @@ from astrai.trainer import (
     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",
@@ -98,7 +71,6 @@ __all__ = [
     "EmbeddingEncoder",
     "EncoderConfig",
     "ExecutorFactory",
-    "GenerationRequest",
     "InferenceEngine",
     "LoRAConfig",
     "Pipeline",
@@ -124,3 +96,5 @@ __all__ = [
     "setup_logging",
     "spawn_parallel_fn",
 ]
+
+setup_logging()

astrai/config/model_config.py

@@ -63,6 +63,11 @@ class AutoRegressiveLMConfig(BaseModelConfig):
         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
@@ -87,6 +92,12 @@ class AutoRegressiveLMConfig(BaseModelConfig):
     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:
@@ -102,6 +113,12 @@ class AutoRegressiveLMConfig(BaseModelConfig):
             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")

astrai/config/train_config.py

@@ -48,6 +48,7 @@ class TrainConfig(BaseConfig):
         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.
+        persistent_workers (bool): Keep DataLoader workers alive between epochs. Defaults to False.
         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.
@@ -61,6 +62,7 @@ class TrainConfig(BaseConfig):
         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.
@@ -97,6 +99,7 @@ class TrainConfig(BaseConfig):
     random_seed: int = 3407
     num_workers: int = 0
     prefetch_factor: Optional[int] = None
+    persistent_workers: bool = False
     pin_memory: bool = False
     collate_fn: Optional[Callable[[List[Any]], Any]] = None
 
@@ -112,6 +115,7 @@ class TrainConfig(BaseConfig):
     val_split: Optional[float] = 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
@@ -187,7 +191,9 @@ class TrainConfig(BaseConfig):
             raise ValueError(f"rollout_top_p must be in (0, 1], got {v}")
         return v
 
-    @field_validator("rollout_top_k", "num_workers", "neftune_alpha")
+    @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}")

astrai/dataset/dataset.py

@@ -25,20 +25,50 @@ function (pure ``record -> Dict[str, Tensor]``) is forwarded to
 
 from abc import ABC, abstractmethod
 from functools import partial
+from pathlib import Path
 from typing import Callable, Dict, List, Optional
 
 import torch
 from torch import Tensor
 from torch.utils.data import Dataset
 
+from astrai.config.preprocess_config import PipelineConfig
 from astrai.dataset.storage import (
     Store,
     StoreFactory,
     detect_format,
 )
 from astrai.factory import BaseFactory
+from astrai.preprocessing.transform import TokenizeTransform
 from astrai.tokenize import AutoTokenizer
 
+_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 _build_jsonl_transform(
+    path: str, tokenizer_path: Optional[str] = None
+) -> Optional["TokenizeTransform"]:
+    """Auto-build a TokenizeTransform for JSONL eager loading.
+
+    Reads ``dataset_config.json`` from the data dir if present, or
+    falls back to the built-in chatml SFT config when *tokenizer_path*
+    is provided.
+    """
+    root = Path(path)
+    config_path = root / "dataset_config.json" if root.is_dir() else None
+    if config_path is not None and config_path.exists():
+        return TokenizeTransform.from_config_file(str(config_path))
+    if tokenizer_path:
+        config = PipelineConfig.from_dict(_DEFAULT_MESSAGES_CONFIG)
+        return TokenizeTransform(config, tokenizer_path)
+    return None
+
 
 def dpo_tokenize(
     record: dict,
@@ -349,16 +379,18 @@ class DatasetFactory(BaseFactory["BaseDataset"]):
         )
         if processor is not None:
             store.load(load_path, processor=processor, **kwargs)
+        elif storage_type == "jsonl":
+            transform = _build_jsonl_transform(load_path, tokenizer_path)
+            if transform is None:
+                raise FileNotFoundError(
+                    f"JSONL dataset config not found. Expected "
+                    f"dataset_config.json alongside *.jsonl files, pass "
+                    f"tokenizer_path= for the built-in messages config, or "
+                    f"use processor= for lazy on-the-fly tokenisation."
+                )
+            store.load(load_path, transform=transform, **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)
+            store.load(load_path, **kwargs)
 
         return cls.create(train_type, store=store)
 

astrai/dataset/storage.py

@@ -55,9 +55,7 @@ from typing import Callable, Dict, List, Optional, Tuple, Union
 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,
@@ -536,19 +534,8 @@ class JsonlStore(Store, Streamable, Recordable):
       ``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,
@@ -569,22 +556,10 @@ class JsonlStore(Store, Streamable, Recordable):
             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)
+            raise ValueError(
+                "JsonlStore eager mode requires transform=. "
+                "Use DatasetFactory.load() which auto-constructs it."
+            )
 
         transformed = transform.apply(records)
         self._normalize(transformed)

astrai/extension/__init__.py

@@ -15,28 +15,34 @@ 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 (
+from astrai.extension.backend import (
     ATTN_BACKEND,
     AttentionBackend,
+    AttentionBackendFactory,
     CudaBackend,
+    FlashAttnBackend,
     TorchNativeBackend,
+    apply_rotary_emb,
     attention,
     attn_backend,
     get_backend,
 )
-from astrai.extension.attention_ops import (
+from astrai.extension.loader import KERNEL_NAMES, is_available
+from astrai.extension.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",

astrai/extension/attention_backend.py

@@ -1,422 +0,0 @@
-"""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 ``TorchNativeBackend`` singleton.
-
-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
-from abc import ABC, abstractmethod
-from contextlib import contextmanager
-from typing import Optional, Union
-
-import torch
-import torch.nn.functional as F
-from torch import Tensor
-
-from astrai.extension.attention_ops import attn_paged_decode, attn_prefill
-from astrai.extension.loader import is_available
-from astrai.inference.core.cache import KVCache
-
-_current_backend: contextvars.ContextVar["AttentionBackend"] = contextvars.ContextVar(
-    "attn_backend"
-)
-
-
-class ATTN_BACKEND(enum.Enum):
-    """Backend selector enum, mirroring ``torch.nn.attention.SDPBackend``."""
-
-    TORCH_NATIVE = "torch_native"
-    CUDA = "cuda"
-
-
-def get_backend() -> "AttentionBackend":
-    """Return the active backend for the current thread/context.
-
-    Falls back to a ``TorchNativeBackend`` singleton when no backend
-    has been activated via ``with``.
-    """
-    try:
-        return _current_backend.get()
-    except LookupError:
-        return _default_backend
-
-
-@contextmanager
-def attn_backend(backend: Union[ATTN_BACKEND, "AttentionBackend", type]):
-    """Context manager to select an attention backend.
-
-    Mirrors ``torch.nn.attention.sdpa_kernel``. Accepts an
-    ``ATTN_BACKEND`` enum value, a backend class, or a backend instance.
-
-    Examples::
-
-        with attn_backend(ATTN_BACKEND.TORCH_NATIVE):
-            ...
-        with attn_backend(TorchNativeBackend):
-            ...
-        with attn_backend(TorchNativeBackend()):
-            ...
-    """
-    if isinstance(backend, ATTN_BACKEND):
-        instance = _BACKEND_REGISTRY[backend]()
-    elif isinstance(backend, type) and issubclass(backend, AttentionBackend):
-        instance = backend()
-    elif isinstance(backend, AttentionBackend):
-        instance = backend
-    else:
-        raise TypeError(
-            f"expected ATTN_BACKEND, AttentionBackend type, 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 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()
-    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."""
-
-
-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.
-    """
-
-    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:
-            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
-            if kv_cache.page_table is not None:
-                indices = kv_cache.page_table
-            else:
-                indices = kv_cache.req_to_token[kv_cache.req_pool_indices, :max_len]
-            if kv_cache.decode_mask is not None:
-                pos_mask = kv_cache.decode_mask
-            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))
-            k = kv_cache.k_buffer[layer_id, indices]
-            v = kv_cache.v_buffer[layer_id, indices]
-
-        n_rep = q.size(2) // k.size(2)
-        if n_rep > 1:
-            k = repeat_kv(k, n_rep)
-            v = repeat_kv(v, n_rep)
-
-        q = q.permute(0, 2, 1, 3)
-        k = k.permute(0, 2, 1, 3)
-        v = v.permute(0, 2, 1, 3)
-
-        out = F.scaled_dot_product_attention(q, k, v, attn_mask, is_causal=is_causal)
-        out = out.permute(0, 2, 1, 3).contiguous().flatten(2)
-        return out
-
-
-_default_backend = TorchNativeBackend()
-
-
-class CudaBackend(AttentionBackend):
-    """CUDA kernel backend with direct KV cache access.
-
-    Decode path: writes K/V to cache, then calls ``attn_paged_decode``
-    with ``page_size=1`` (each token slot is a single-token "page").
-    The ``req_to_token`` table serves directly as the page table.
-
-    Prefill path: writes K/V to cache, gathers full-sequence K/V via
-    indirect indexing (same as TorchNativeBackend), then calls
-    ``attn_prefill``.
-
-    Training path (``kv_cache is None``): calls ``attn_prefill`` directly
-    on the projected q/k/v.
-
-    Falls back to ``TorchNativeBackend`` for any path where the
-    corresponding CUDA kernel is not available.
-    """
-
-    def __init__(self):
-        self._fallback = TorchNativeBackend()
-
-    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 or not is_available("attn_paged_decode"):
-            return self._fallback.fwd_decode(
-                q, k, v, kv_cache, layer_id, attn_mask, is_causal
-            )
-
-        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
-
-        if kv_cache.page_table is not None:
-            page_table = kv_cache.page_table
-        else:
-            page_table = kv_cache.req_to_token[kv_cache.req_pool_indices, :max_len]
-
-        k_cache = kv_cache.k_buffer[layer_id].unsqueeze(1)
-        v_cache = kv_cache.v_buffer[layer_id].unsqueeze(1)
-
-        if q.size(0) == 1:
-            mask = None
-        elif kv_cache.decode_mask is not None:
-            mask = kv_cache.decode_mask
-        else:
-            mask = (
-                torch.arange(max_len, device=q.device)[None, :]
-                < kv_cache.seq_lens[:, None]
-            )
-
-        out = attn_paged_decode(
-            q,
-            page_table,
-            k_cache,
-            v_cache,
-            page_size=1,
-            kv_len=max_len,
-            mask=mask,
-            is_causal=is_causal,
-        )
-
-        out = out.flatten(2)
-        return out
-
-    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:
-            if is_available("attn_prefill"):
-                out = attn_prefill(q, k, v, mask=attn_mask, is_causal=is_causal)
-                return out.flatten(2)
-            return self._fallback.fwd_prefill(
-                q, k, v, kv_cache, layer_id, attn_mask, is_causal
-            )
-
-        if not is_available("attn_prefill"):
-            return self._fallback.fwd_prefill(
-                q, k, v, kv_cache, layer_id, attn_mask, is_causal
-            )
-
-        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]
-        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))
-        k_full = kv_cache.k_buffer[layer_id, indices]
-        v_full = kv_cache.v_buffer[layer_id, indices]
-
-        out = attn_prefill(q, k_full, v_full, mask=attn_mask, is_causal=is_causal)
-        return out.flatten(2)
-
-
-_BACKEND_REGISTRY: dict[ATTN_BACKEND, type[AttentionBackend]] = {
-    ATTN_BACKEND.TORCH_NATIVE: TorchNativeBackend,
-    ATTN_BACKEND.CUDA: CudaBackend,
-}

astrai/extension/attention_ops.py

@@ -1,117 +0,0 @@
-"""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 torch
-
-from astrai.extension.loader import _available, _modules
-
-
-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: torch.Tensor | None = 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=1
-    )
-
-
-def attn_prefill(
-    q: torch.Tensor,
-    k: torch.Tensor,
-    v: torch.Tensor,
-    mask: torch.Tensor | None = 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=1
-    )
-
-
-def attn_paged_decode(
-    q: torch.Tensor,
-    page_table: torch.Tensor,
-    k_cache: torch.Tensor,
-    v_cache: torch.Tensor,
-    page_size: int,
-    kv_len: int,
-    mask: torch.Tensor | None = None,
-    is_causal: bool = False,
-) -> torch.Tensor:
-    """Paged GQA decode attention (q_len == 1, direct page-table access).
-
-    Args:
-        q: [batch, 1, n_heads, head_dim] (blhd, bf16)
-        page_table: [batch, max_pages] (int64)
-        k_cache: [n_pages, page_size, n_kv_heads, head_dim] (bf16)
-        v_cache: same as k_cache
-        page_size: tokens per page
-        kv_len: actual sequence length per request
-        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_paged_decode")
-    causal_offset = (kv_len - 1) if is_causal else -1
-    return _modules["attn_paged_decode"].attn_paged_decode(
-        q,
-        page_table,
-        k_cache,
-        v_cache,
-        page_size,
-        kv_len,
-        mask=mask,
-        causal_offset=causal_offset,
-        layout=1,
-    )

astrai/extension/backend/__init__.py

@@ -0,0 +1,27 @@
+"""Backend selection, fallbacks, and execution policies."""
+
+from astrai.extension.backend.attention import (
+    ATTN_BACKEND,
+    AttentionBackend,
+    AttentionBackendFactory,
+    CudaBackend,
+    FlashAttnBackend,
+    TorchNativeBackend,
+    attention,
+    attn_backend,
+    get_backend,
+)
+from astrai.extension.backend.rotary import apply_rotary_emb
+
+__all__ = [
+    "ATTN_BACKEND",
+    "AttentionBackend",
+    "AttentionBackendFactory",
+    "CudaBackend",
+    "FlashAttnBackend",
+    "TorchNativeBackend",
+    "apply_rotary_emb",
+    "attention",
+    "attn_backend",
+    "get_backend",
+]

astrai/extension/backend/attention.py

@@ -0,0 +1,702 @@
+"""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 os
+import threading
+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.loader import is_available
+from astrai.extension.ops.attention import (
+    attn_paged_decode,
+    attn_paged_prefill,
+)
+from astrai.factory import BaseFactory
+
+try:
+    import flash_attn as _flash_attn
+except Exception:
+    _flash_attn = None
+
+if TYPE_CHECKING:
+    from astrai.inference.cache import KVCache
+
+
+_default_backend: Optional["AttentionBackend"] = None
+_default_backend_lock = threading.Lock()
+_env_backend_name: Optional[str] = None
+_env_backend: Optional["AttentionBackend"] = None
+_current_backend: contextvars.ContextVar[Optional["AttentionBackend"]] = (
+    contextvars.ContextVar("attn_backend", default=None)
+)
+
+
+@functools.lru_cache(maxsize=1)
+def flash_attn_available() -> bool:
+    if not torch.cuda.is_available():
+        return False
+    fa = _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
+
+
+class ATTN_BACKEND(enum.Enum):
+    """Backend selector enum, mirroring ``torch.nn.attention.SDPBackend``."""
+
+    TORCH_NATIVE = "torch_native"
+    CUDA = "cuda"
+    FLASH = "flash"
+
+
+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,
+    fwd: Optional[str],
+) -> 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 (
+            fwd in ("prefill", "decode")
+            and kv_cache is not None
+            and q.ndim == 3
+            and q.dtype == torch.bfloat16
+            and q.size(-1) in (32, 64, 128, 256)
+            and is_available(f"attn_paged_{fwd}")
+        )
+    if isinstance(backend, FlashAttnBackend):
+        if not flash_attn_available():
+            return False
+        if q.dtype not in (torch.float16, torch.bfloat16):
+            return False
+        if fwd is not None:
+            return q.ndim == 3 and hasattr(_flash_attn, "flash_attn_varlen_func")
+        if attn_mask is None or is_causal:
+            return True
+        return attn_mask.dim() == 4
+    return True
+
+
+def _resolve_default_backend() -> "AttentionBackend":
+    """Pick the highest-priority available backend (cuda -> flash -> torch).
+
+    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.
+    """
+    return _priority_backends()[0]
+
+
+def _environment_backend() -> Optional["AttentionBackend"]:
+    """Resolve the process-wide ``ASTR_BACKEND`` override, if configured."""
+    global _env_backend, _env_backend_name
+    name = os.environ.get("ASTR_BACKEND", "").strip().lower()
+    if not name:
+        return None
+    if name != _env_backend_name:
+        with _default_backend_lock:
+            if name != _env_backend_name:
+                try:
+                    _env_backend = AttentionBackendFactory.create(name)
+                except (ValueError, RuntimeError):
+                    _env_backend = None
+                _env_backend_name = name
+    return _env_backend
+
+
+def _resolve_backend(
+    backend: Optional[Union[str, ATTN_BACKEND, "AttentionBackend", type]] = None,
+) -> "AttentionBackend":
+    """Resolve a backend configuration, defaulting to the process policy."""
+    if backend is not None:
+        if isinstance(backend, ATTN_BACKEND):
+            return AttentionBackendFactory.create(backend.value)
+        if isinstance(backend, str):
+            return AttentionBackendFactory.create(backend)
+        if isinstance(backend, type) and issubclass(backend, AttentionBackend):
+            return backend()
+        if isinstance(backend, AttentionBackend):
+            return backend
+        raise TypeError(
+            f"expected a registered name, ATTN_BACKEND, AttentionBackend type, "
+            f"or instance, got {type(backend).__name__}"
+        )
+
+    global _default_backend
+    if _default_backend is None:
+        with _default_backend_lock:
+            if _default_backend is None:
+                _default_backend = _resolve_default_backend()
+    return _default_backend
+
+
+def get_backend(
+    use_default: bool = True,
+) -> Optional["AttentionBackend"]:
+    """Return the context override, optionally falling back to the process default.
+
+    ``ASTR_BACKEND`` is a process-wide override and takes precedence over the
+    context value. Pass ``use_default=False`` at request submission to retain
+    only an environment override or the caller's :func:`attn_backend` value.
+    """
+    return (
+        _environment_backend()
+        or _current_backend.get()
+        or (_resolve_backend() if use_default else None)
+    )
+
+
+@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()):
+            ...
+    """
+    instance = _resolve_backend(backend)
+    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."""
+    if n_rep == 1:
+        return x
+    n_heads, head_dim = x.shape[-2:]
+    return (
+        x.unsqueeze(-2)
+        .expand(*x.shape[:-2], n_heads, n_rep, head_dim)
+        .reshape(*x.shape[:-2], n_heads * n_rep, head_dim)
+    )
+
+
+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,
+    fwd: Optional[str] = None,
+) -> 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]
+    """
+    explicit = get_backend(use_default=False)
+    backend = get_backend()
+    if fwd is None and explicit is None:
+        backend = TorchNativeBackend()
+    if not _backend_supports(backend, q, kv_cache, attn_mask, is_causal, fwd):
+        if explicit is not None:
+            raise RuntimeError(
+                f"Explicitly-set backend {type(backend).__name__} cannot "
+                f"handle this attention call (shape={q.shape}, "
+                f"dtype={q.dtype}, kv_cache={'none' if kv_cache is None else 'present'}, "
+                f"attn_mask={'none' if attn_mask is None else 'present'}). "
+                f"Remove the attn_backend() context or switch to a compatible backend."
+            )
+        for candidate in _priority_backends():
+            if isinstance(candidate, type(backend)):
+                continue
+            if _backend_supports(candidate, q, kv_cache, attn_mask, is_causal, fwd):
+                backend = candidate
+                break
+    return backend.forward(q, k, v, kv_cache, layer_id, attn_mask, is_causal, fwd)
+
+
+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,
+        fwd: Optional[str] = None,
+    ) -> 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 fwd == "decode":
+            return self.fwd_decode(q, k, v, kv_cache, layer_id, attn_mask, is_causal)
+        if fwd == "prefill" or fwd is None:
+            return self.fwd_prefill(q, k, v, kv_cache, layer_id, attn_mask, is_causal)
+        raise ValueError(f"unsupported attention forward mode: {fwd}")
+
+    @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 q.ndim == 4:
+            n_rep = q.size(2) // k.size(2)
+            if n_rep > 1:
+                k = repeat_kv(k, n_rep)
+                v = repeat_kv(v, n_rep)
+            return (
+                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,
+                )
+                .permute(0, 2, 1, 3)
+                .contiguous()
+            )
+
+        if kv_cache is None or kv_cache.qo_indptr is None:
+            raise ValueError("packed attention requires KV cache metadata")
+        kv_cache.k_buffer[layer_id, kv_cache.out_cache_loc] = k
+        kv_cache.v_buffer[layer_id, kv_cache.out_cache_loc] = v
+        outputs = []
+        n_rep = q.size(1) // k.size(1)
+        for i in range(kv_cache.req_pool_indices.numel()):
+            q_start = int(kv_cache.qo_indptr[i])
+            q_end = int(kv_cache.qo_indptr[i + 1])
+            indices = kv_cache.req_to_token[
+                kv_cache.req_pool_indices[i], : kv_cache.seq_lens[i]
+            ]
+            k_i = kv_cache.k_buffer[layer_id, indices]
+            v_i = kv_cache.v_buffer[layer_id, indices]
+            if n_rep > 1:
+                k_i = repeat_kv(k_i, n_rep)
+                v_i = repeat_kv(v_i, n_rep)
+            q_len = q_end - q_start
+            kv_len = k_i.size(0)
+            q_pos = torch.arange(kv_len - q_len, kv_len, device=q.device)
+            causal_mask = q_pos[:, None] >= torch.arange(kv_len, device=q.device)
+            out = F.scaled_dot_product_attention(
+                q[q_start:q_end].transpose(0, 1).unsqueeze(0),
+                k_i.transpose(0, 1).unsqueeze(0),
+                v_i.transpose(0, 1).unsqueeze(0),
+                attn_mask=causal_mask,
+            )
+            outputs.append(out.squeeze(0).transpose(0, 1))
+        return torch.cat(outputs)
+
+
+@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)")
+
+        kv_indptr = kv_cache.kv_indptr
+
+        out = attn_paged_decode(
+            q,
+            kv_cache.k_buffer[layer_id],
+            kv_cache.v_buffer[layer_id],
+            kv_cache.req_to_token,
+            kv_cache.req_pool_indices,
+            kv_indptr,
+            new_k=k,
+            new_v=v,
+            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
+
+    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
+        kv_cache.k_buffer[layer_id, loc] = k
+        kv_cache.v_buffer[layer_id, loc] = v
+
+        out = attn_paged_prefill(
+            q,
+            kv_cache.k_buffer[layer_id],
+            kv_cache.v_buffer[layer_id],
+            kv_cache.req_to_token,
+            kv_cache.req_pool_indices,
+            kv_cache.kv_indptr,
+            kv_cache.qo_indptr,
+            kv_cache.q_tile_to_batch,
+            kv_cache.q_tile_to_index,
+            attn_mask,
+            is_causal=is_causal,
+        )
+        return out
+
+
+@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_packed(q, k, v, kv_cache, layer_id)
+
+    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 q.ndim == 3:
+            return self._forward_packed(q, k, v, kv_cache, layer_id)
+        return self._forward_dense(q, k, v, attn_mask, is_causal)
+
+    def _forward_dense(
+        self,
+        q: Tensor,
+        k: Tensor,
+        v: Tensor,
+        attn_mask: Optional[Tensor] = None,
+        is_causal: bool = False,
+    ) -> Tensor:
+        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 and attn_mask.dim() != 4:
+            raise ValueError(
+                "FlashAttnBackend does not support a custom attention mask; "
+                "use a causal mask or select TorchNativeBackend."
+            )
+        fa = _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 or (attn_mask is not None and attn_mask.dim() == 4),
+        )
+        return out.contiguous()
+
+    def _forward_packed(
+        self,
+        q: Tensor,
+        k: Tensor,
+        v: Tensor,
+        kv_cache: "KVCache",
+        layer_id: int,
+    ) -> Tensor:
+        fa = _flash_attn
+        if fa is None or not hasattr(fa, "flash_attn_varlen_func"):
+            raise RuntimeError("packed inference requires flash_attn_varlen_func")
+        kv_cache.k_buffer[layer_id, kv_cache.out_cache_loc] = k
+        kv_cache.v_buffer[layer_id, kv_cache.out_cache_loc] = v
+        page_table = kv_cache.req_to_token[
+            kv_cache.req_pool_indices, : kv_cache.max_len
+        ]
+        positions = torch.arange(kv_cache.max_len, device=q.device)
+        indices = page_table[positions.unsqueeze(0) < kv_cache.seq_lens.unsqueeze(1)]
+        k_flat = kv_cache.k_buffer[layer_id, indices].contiguous()
+        v_flat = kv_cache.v_buffer[layer_id, indices].contiguous()
+        out = fa.flash_attn_varlen_func(
+            q.contiguous(),
+            k_flat,
+            v_flat,
+            kv_cache.qo_indptr,
+            kv_cache.kv_indptr,
+            int((kv_cache.qo_indptr[1:] - kv_cache.qo_indptr[:-1]).max()),
+            int(kv_cache.seq_lens.max()),
+            dropout_p=0.0,
+            causal=True,
+        )
+        return out

astrai/extension/backend/rotary.py

@@ -11,6 +11,7 @@ import torch
 from torch import Tensor
 
 from astrai.extension.loader import is_available
+from astrai.extension.ops.rotary import rotary_emb as _cuda_rotary
 
 _cache = {"available": None}
 
@@ -26,7 +27,7 @@ def _torch_apply(x: 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_complex = torch.complex(cos, sin).unsqueeze(2)
+    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)
@@ -48,7 +49,5 @@ def apply_rotary_emb(x: Tensor, freqs_cis: Tensor) -> Tensor:
         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)

astrai/extension/fp8.py

@@ -0,0 +1,334 @@
+"""FP8 training: scaling state and aten::linear dispatch.
+
+Layered (see also ``ops/fp8.py`` for the CUDA interface adapter):
+
+1. Kernel interface: ``ops.fp8`` - the only module touching the pybind.
+2. Training state (this module): per-tensor scales, amax history, delayed
+   scaling, and the ``fp8_autocast`` context (TE-style, like
+   ``torch.autocast``).
+3. aten::linear integration (this module): registers the CUDA impl and the
+   M/N alignment guard.
+
+Usage::
+
+    from astrai.extension.fp8 import fp8_autocast
+
+    with fp8_autocast(enabled=True):
+        logits = model(input_ids)
+        loss.backward()
+
+Importing this module registers the aten::linear CUDA implementation.
+"""
+
+from contextlib import contextmanager
+
+import torch
+from torch.library import Library
+
+from astrai.extension.ops.fp8 import (
+    linear_backward_scaled,
+    linear_forward_scaled,
+)
+
+E4M3_MAX = 448.0
+
+
+# ---------------------------------------------------------------------------
+# Layer 2: training state (scales, amax history, delayed scaling, autocast)
+# ---------------------------------------------------------------------------
+
+
+class FP8TensorMeta:
+    """Scales + amax state for one weight tensor and its paired activations.
+
+    - weight: delayed scale from a 16-step amax history window (TE style)
+    - x/g: delayed one step, reuse the quantize kernel's free atomic amax
+    """
+
+    __slots__ = (
+        "scale",
+        "scale_inv",
+        "amax_history",
+        "idx",
+        "x_scale",
+        "x_scale_inv",
+        "x_history",
+        "x_idx",
+        "g_scale",
+        "g_scale_inv",
+        "g_history",
+        "g_idx",
+        "w_init",
+        "x_init",
+        "g_init",
+    )
+
+    def __init__(self, device: torch.device, update_interval: int):
+        self.scale = torch.ones(1, device=device, dtype=torch.float32)
+        self.scale_inv = torch.ones(1, device=device, dtype=torch.float32)
+        self.amax_history = torch.ones(
+            update_interval, device=device, dtype=torch.float32
+        )
+        self.idx = 0
+        self.x_scale = torch.ones(1, device=device, dtype=torch.float32)
+        self.x_scale_inv = torch.ones(1, device=device, dtype=torch.float32)
+        self.x_history = torch.ones(update_interval, device=device, dtype=torch.float32)
+        self.x_idx = 0
+        self.g_scale = torch.ones(1, device=device, dtype=torch.float32)
+        self.g_scale_inv = torch.ones(1, device=device, dtype=torch.float32)
+        self.g_history = torch.ones(update_interval, device=device, dtype=torch.float32)
+        self.g_idx = 0
+        self.w_init = False
+        self.x_init = False
+        self.g_init = False
+
+    def init_scale(self, t: torch.Tensor) -> None:
+        """Immediate scale from the current amax; used on the first call.
+
+        A scale of 1 would underflow small activations/gradients (e4m3 min
+        normal is 2^-6); initialize from the actual amax once, then delayed
+        updates take over.
+        """
+        amax = t.abs().amax().to(torch.float32).clamp_min(1e-12)
+        self.scale.copy_(amax / E4M3_MAX)
+        self.scale_inv.copy_(E4M3_MAX / amax)
+        self.record(amax)
+
+    def push_x_scale(self, amax: torch.Tensor) -> None:
+        """Window update for the activation scale (delayed, TE style)."""
+        self.x_history[self.x_idx] = amax.reshape(())
+        self.x_idx = (self.x_idx + 1) % self.x_history.numel()
+        m = self.x_history.max()
+        self.x_scale.copy_(m / E4M3_MAX)
+        self.x_scale_inv.copy_(E4M3_MAX / m)
+
+    def push_g_scale(self, amax: torch.Tensor) -> None:
+        """Window update for the gradient scale (delayed, TE style)."""
+        self.g_history[self.g_idx] = amax.reshape(())
+        self.g_idx = (self.g_idx + 1) % self.g_history.numel()
+        m = self.g_history.max()
+        self.g_scale.copy_(m / E4M3_MAX)
+        self.g_scale_inv.copy_(E4M3_MAX / m)
+
+    def record(self, amax: torch.Tensor) -> None:
+        """Push the latest amax into the ring buffer (device-side copy, no sync)."""
+        self.amax_history[self.idx] = amax.reshape(())
+        self.idx = (self.idx + 1) % self.amax_history.numel()
+
+    def refresh(self) -> None:
+        """Recompute scale from the amax history window (delayed scaling)."""
+        amax = self.amax_history.max()
+        if amax > 0:
+            self.scale.copy_(amax / E4M3_MAX)
+            self.scale_inv.copy_(E4M3_MAX / amax)
+
+
+class FP8State:
+    """Global fp8 training state, TE-style."""
+
+    def __init__(self, update_interval: int = 16):
+        self.enabled = False
+        self.update_interval = update_interval
+        self.step_count = 0
+        self._metas: dict[tuple, FP8TensorMeta] = {}
+        self._last_device: torch.device | None = None
+
+    def _get_device(self, t: torch.Tensor) -> torch.device:
+        if self._last_device is None:
+            self._last_device = t.device
+        return t.device
+
+    def get_weight_meta(self, w: torch.Tensor) -> FP8TensorMeta:
+        key = (w.data_ptr(), w.shape, w.dtype)
+        meta = self._metas.get(key)
+        if meta is None:
+            meta = FP8TensorMeta(self._get_device(w), self.update_interval)
+            self._metas[key] = meta
+        return meta
+
+    def step(self) -> None:
+        """Advance the counter and refresh all weight scales every N steps."""
+        self.step_count += 1
+        if self.step_count % self.update_interval == 0:
+            for meta in self._metas.values():
+                meta.refresh()
+
+    def reset(self) -> None:
+        self.enabled = False
+        self.step_count = 0
+        self._metas.clear()
+        self._last_device = None
+
+
+# Global singleton: autograd backward runs on the engine worker threads, so
+# thread-local state would lose the fp8 flag during loss.backward(). The GIL
+# protects Python-side mutation; the CUDA kernels take their own mutex.
+_state = FP8State()
+
+
+def fp8_state() -> FP8State:
+    return _state
+
+
+@contextmanager
+def fp8_autocast(enabled: bool = True, update_interval: int = 16):
+    """Autocast-style context: fp8 linear dispatch on this thread.
+
+    Usage::
+
+        with fp8_autocast(enabled=True):
+            logits = model(input_ids)   # aten::linear -> fp8 path
+            loss.backward()
+
+    The scale-update counter advances once per ``enter`` (one training step),
+    refreshing weight scales from their amax history every ``update_interval``.
+    """
+    state = fp8_state()
+    prev_enabled = state.enabled
+    prev_interval = state.update_interval
+    state.enabled = enabled
+    state.update_interval = update_interval
+    try:
+        if enabled:
+            state.step()
+        yield
+    finally:
+        state.enabled = prev_enabled
+        state.update_interval = prev_interval
+
+
+def fp8_linear_forward(x: torch.Tensor, w: torch.Tensor, bias=None):
+    """TE-style scaled fp8 linear forward (called from the aten::linear impl).
+
+    x uses the delayed scale of its paired weight meta (amax from the previous
+    forward of this linear); the quantize kernel emits the current amax for the
+    next step. No extra abs/max reduce.
+    """
+    if bias is None:
+        bias = torch.empty(0, device=x.device, dtype=x.dtype)
+    state = fp8_state()
+    meta = state.get_weight_meta(w)
+    if not meta.w_init:
+        meta.init_scale(w)
+        meta.w_init = True
+    if not meta.x_init:
+        amax = x.abs().amax().to(torch.float32).clamp_min(1e-12)
+        meta.x_history.fill_(amax)
+        meta.x_scale.copy_(amax / E4M3_MAX)
+        meta.x_scale_inv.copy_(E4M3_MAX / amax)
+        meta.x_init = True
+    amax_x = torch.empty(1, device=x.device, dtype=torch.float32)
+    amax_w = torch.empty(1, device=x.device, dtype=torch.float32)
+    out = linear_forward_scaled(
+        x,
+        w,
+        bias,
+        meta.x_scale,
+        meta.scale,
+        meta.x_scale_inv,
+        meta.scale_inv,
+        amax_x,
+        amax_w,
+    )
+    meta.record(amax_w)
+    meta.push_x_scale(amax_x)
+    return out
+
+
+def fp8_linear_backward(g, x, w, masks):
+    """TE-style scaled fp8 linear backward (called from aten::linear_backward)."""
+    state = fp8_state()
+    meta = state.get_weight_meta(w)
+    if not meta.g_init:
+        amax = g.abs().amax().to(torch.float32).clamp_min(1e-12)
+        meta.g_history.fill_(amax)
+        meta.g_scale.copy_(amax / E4M3_MAX)
+        meta.g_scale_inv.copy_(E4M3_MAX / amax)
+        meta.g_init = True
+    amax_g = torch.empty(1, device=g.device, dtype=torch.float32)
+    out = linear_backward_scaled(
+        g,
+        x,
+        w,
+        masks,
+        meta.g_scale,
+        meta.scale,
+        meta.x_scale,
+        meta.g_scale_inv,
+        meta.scale_inv,
+        meta.x_scale_inv,
+        amax_g,
+    )
+    meta.push_g_scale(amax_g)
+    return out
+
+
+# ---------------------------------------------------------------------------
+# Layer 3: aten::linear integration
+# ---------------------------------------------------------------------------
+
+
+def fp8_linear_enable(enabled: bool = True) -> None:
+    """Toggle fp8 dispatch for aten::linear (global; backward runs on engine
+    worker threads, so a thread-local flag would be lost during backward)."""
+    fp8_state().enabled = enabled
+
+
+def fp8_linear_enabled() -> bool:
+    return fp8_state().enabled
+
+
+def _fp8_supported(x: torch.Tensor, w: torch.Tensor) -> bool:
+    """cuBLASLt fp8 requires M % 16 == 0 and N % 16 == 0 (K is padded)."""
+    m = x.numel() // x.size(-1)
+    return m % 16 == 0 and w.size(0) % 16 == 0
+
+
+def _linear_cuda_impl(x: torch.Tensor, w: torch.Tensor, bias=None):
+    if (
+        fp8_linear_enabled()
+        and x.dtype == torch.bfloat16
+        and w.dtype == torch.bfloat16
+        and _fp8_supported(x, w)
+    ):
+        return fp8_linear_forward(x, w, bias)
+    return torch.ops.aten.linear.default.redispatch(
+        torch._C.DispatchKeySet(torch._C.DispatchKey.CompositeImplicitAutograd),
+        x,
+        w,
+        bias,
+    )
+
+
+def _linear_backward_cuda_impl(input_tensor, grad_output, weight, output_mask):
+    if (
+        fp8_linear_enabled()
+        and weight.dtype == torch.bfloat16
+        and _fp8_supported(grad_output, weight)
+    ):
+        return fp8_linear_backward(grad_output, input_tensor, weight, list(output_mask))
+    compute_dtype = weight.dtype
+    grad = grad_output.to(compute_dtype)
+    grad_2d = grad.reshape(-1, weight.size(0))
+    input_2d = input_tensor.reshape(-1, input_tensor.size(-1)).to(compute_dtype)
+    grad_input = (
+        torch.mm(grad_2d, weight)
+        if output_mask[0]
+        else torch.empty(0, device=input_tensor.device, dtype=input_tensor.dtype)
+    )
+    grad_weight = (
+        torch.mm(grad_2d.t(), input_2d)
+        if output_mask[1]
+        else torch.empty(0, device=input_tensor.device, dtype=input_tensor.dtype)
+    )
+    grad_bias = (
+        grad.sum(dim=0)
+        if output_mask[2]
+        else torch.empty(0, device=input_tensor.device, dtype=input_tensor.dtype)
+    )
+    return grad_input.reshape_as(input_tensor), grad_weight, grad_bias
+
+
+_lib = Library("aten", "IMPL", "CUDA")
+_lib.impl("linear", _linear_cuda_impl)
+_lib.impl("linear_backward", _linear_backward_cuda_impl)

astrai/extension/loader.py

@@ -11,7 +11,14 @@ import logging
 
 logger = logging.getLogger(__name__)
 
-KERNEL_NAMES = ["attn_decode", "attn_prefill", "attn_paged_decode", "rotary_emb"]
+KERNEL_NAMES = [
+    "attn_decode",
+    "attn_prefill",
+    "attn_paged_decode",
+    "attn_paged_prefill",
+    "rotary_emb",
+    "fp8_mm",
+]
 
 _available: dict[str, bool] = {}
 _modules: dict[str, object] = {}

astrai/extension/ops/__init__.py

@@ -0,0 +1,19 @@
+"""Stateless wrappers around compiled extension kernels."""
+
+from astrai.extension.ops.attention import (
+    TensorLayout,
+    attn_decode,
+    attn_paged_decode,
+    attn_paged_prefill,
+    attn_prefill,
+)
+from astrai.extension.ops.rotary import rotary_emb
+
+__all__ = [
+    "TensorLayout",
+    "attn_decode",
+    "attn_paged_decode",
+    "attn_paged_prefill",
+    "attn_prefill",
+    "rotary_emb",
+]

astrai/extension/ops/attention.py

@@ -0,0 +1,200 @@
+"""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,
+    new_k: Optional[torch.Tensor] = None,
+    new_v: Optional[torch.Tensor] = None,
+    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] (int32) — token -> slot
+        req_pool_indices: [batch] (int32) — rows into req_to_token
+        kv_indptr: [batch+1] (int32) — prefix sum of per-request seq_lens
+        new_k: current-token K to append, [batch, n_kv_heads, head_dim]
+        new_v: current-token V to append, same shape as new_k
+        mask: 2D [batch, max_context_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,
+        new_k=new_k,
+        new_v=new_v,
+        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,
+    q_tile_to_batch: torch.Tensor,
+    q_tile_to_index: torch.Tensor,
+    mask: Optional[torch.Tensor] = None,
+    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] (int32)
+        req_pool_indices: [batch] (int32)
+        kv_indptr: [batch+1] (int32) — prefix sum of per-request kv_lens
+        qo_indptr: [batch+1] (int32) — prefix sum of per-request q_lens
+        q_tile_to_batch: [num_q_tiles] (int32) — request index per Q tile
+        q_tile_to_index: [num_q_tiles] (int32) — local Q tile index per request
+        mask: 4D [batch, 1, q_len, kv_len] (bool, True=keep) or None
+        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,
+        q_tile_to_batch,
+        q_tile_to_index,
+        mask,
+        causal_offset=causal_offset,
+    )

astrai/extension/ops/fp8.py

@@ -0,0 +1,127 @@
+"""FP8 CUDA kernel interface adapter (the only module touching the pybind.
+
+Isolates the ``fp8_mm`` CUDA extension behind stable Python functions:
+- availability / dtype checks and clear errors
+- torch.library ``custom::fp8_mm`` registration (meta + CPU fallback)
+- quantize-in-GEMM primitives used by ``fp8.py`` training state
+
+Policy (scales, amax history, delayed scaling, autocast) lives in ``fp8.py``;
+this module is stateless.
+"""
+
+import torch
+from torch.library import custom_op
+
+from astrai.extension.loader import get_module, is_available
+
+
+def _mod():
+    if not is_available("fp8_mm"):
+        raise RuntimeError(
+            "CUDA kernel 'fp8_mm' is not available. Build with CSRC_KERNELS=true."
+        )
+    return get_module("fp8_mm")
+
+
+@custom_op("custom::fp8_mm", mutates_args=())
+def fp8_mm(
+    a: torch.Tensor, b: torch.Tensor, sx: torch.Tensor, sw: torch.Tensor
+) -> torch.Tensor:
+    """BF16 inputs, fused FP8 GEMM with FP32 accumulation and BF16 output."""
+
+
+@fp8_mm.register_fake
+def _fp8_mm_fake(a, b, sx, sw):
+    return torch.empty((a.size(0), b.size(0)), device=a.device, dtype=torch.bfloat16)
+
+
+@fp8_mm.register_kernel("cuda")
+def _fp8_mm_cuda(a, b, sx, sw):
+    if not (a.dtype == torch.bfloat16 and b.dtype == torch.bfloat16):
+        raise TypeError(f"bf16 GEMM requires bf16 inputs, got {a.dtype}/{b.dtype}")
+    return _mod().fp8_mm(a, b, sx, sw)
+
+
+@fp8_mm.register_kernel("cpu")
+def _fp8_mm_cpu(a, b, sx, sw):
+    return torch.mm(a.float(), b.float().t()).to(torch.bfloat16)
+
+
+@custom_op("custom::fp8_mm_prequant", mutates_args=())
+def fp8_mm_prequant(
+    a: torch.Tensor, b: torch.Tensor, scale: torch.Tensor
+) -> torch.Tensor:
+    """Pre-quantized FP8 inputs, fused FP8 GEMM, FP32 accumulation, BF16 out."""
+
+
+@fp8_mm_prequant.register_fake
+def _fp8_mm_prequant_fake(a, b, scale):
+    return torch.empty((a.size(0), b.size(0)), device=a.device, dtype=torch.bfloat16)
+
+
+@fp8_mm_prequant.register_kernel("cuda")
+def _fp8_mm_prequant_cuda(a, b, scale):
+    if not (a.dtype == torch.float8_e4m3fn and b.dtype == torch.float8_e4m3fn):
+        raise TypeError(
+            f"pre-quantized FP8 GEMM requires fp8 inputs, got {a.dtype}/{b.dtype}"
+        )
+    return _mod().fp8_mm_prequant(a, b, scale)
+
+
+@fp8_mm_prequant.register_kernel("cpu")
+def _fp8_mm_prequant_cpu(a, b, scale):
+    return (a.float() @ b.float().t() * scale).to(torch.bfloat16)
+
+
+@custom_op("custom::fp8_mm_prequant_fp8", mutates_args=())
+def fp8_mm_prequant_fp8(
+    a: torch.Tensor, b: torch.Tensor, scale: torch.Tensor, out_scale: torch.Tensor
+) -> torch.Tensor:
+    """FP8 inputs and FP8 output: fused FP8 GEMM with FP32 accumulation."""
+
+
+@fp8_mm_prequant_fp8.register_fake
+def _fp8_mm_prequant_fp8_fake(a, b, scale, out_scale):
+    return torch.empty((a.size(0), b.size(0)), device=a.device, dtype=a.dtype)
+
+
+@fp8_mm_prequant_fp8.register_kernel("cuda")
+def _fp8_mm_prequant_fp8_cuda(a, b, scale, out_scale):
+    if not (a.dtype == torch.float8_e4m3fn and b.dtype == torch.float8_e4m3fn):
+        raise TypeError(
+            f"pre-quantized FP8 GEMM requires fp8 inputs, got {a.dtype}/{b.dtype}"
+        )
+    return _mod().fp8_mm_prequant_fp8(a, b, scale, out_scale)
+
+
+@fp8_mm_prequant_fp8.register_kernel("cpu")
+def _fp8_mm_prequant_fp8_cpu(a, b, scale, out_scale):
+    return (a.float() @ b.float().t() * scale * out_scale).to(torch.float8_e4m3fn)
+
+
+def linear_forward_scaled(x, w, bias, sx, sw, sx_inv, sw_inv, amax_x, amax_w):
+    """Quantize BF16 inputs to FP8, accumulate in FP32, and return BF16.
+
+    x/w: [..., K] / [N, K] bf16; sx/sw and their inverses control the fused
+    E4M3 conversion; amax_x/amax_w receive the input max-abs values.
+    """
+    if not (x.dtype == torch.bfloat16 and w.dtype == torch.bfloat16):
+        raise TypeError(f"fp8 forward requires bf16 inputs, got {x.dtype}/{w.dtype}")
+    return _mod().fp8_linear_forward_scaled(
+        x, w, bias, sx, sw, sx_inv, sw_inv, amax_x, amax_w
+    )
+
+
+def linear_backward_scaled(g, x, w, masks, sg, sw, sx, sg_inv, sw_inv, sx_inv, amax_g):
+    """dX = g @ W, dW = g^T @ X, dB = sum(g) with per-tensor scales."""
+    if not (
+        g.dtype == torch.bfloat16
+        and x.dtype == torch.bfloat16
+        and w.dtype == torch.bfloat16
+    ):
+        raise TypeError(
+            f"fp8 backward requires bf16 inputs, got {g.dtype}/{x.dtype}/{w.dtype}"
+        )
+    return _mod().fp8_linear_backward_scaled(
+        g, x, w, masks, sg, sw, sx, sg_inv, sw_inv, sx_inv, amax_g
+    )

astrai/extension/ops/rotary.py

@@ -2,10 +2,10 @@
 
 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``.
+responsibility of ``astrai.extension.backend.rotary.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.
+Layout: x is packed [tokens, n_heads, head_dim] or dense
+[batch, seq_len, n_heads, head_dim]. ``freqs_cis`` has matching token axes.
 """
 
 import torch
@@ -25,11 +25,11 @@ 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
+        x: packed 3D or dense 4D bf16 tensor.
+        freqs_cis: matching token axes followed by [head_dim/2, 2].
 
     Returns:
-        [batch, seq_len, n_heads, head_dim] (bf16)
+        Tensor with the same shape as ``x``.
     """
     _check_available()
     if not x.is_contiguous():

astrai/inference/__init__.py

@@ -1,15 +1,29 @@
 """Inference module for continuous batching.
 
-Layers:
-  - core/:        Core inference loop (cache, executor, scheduler, task)
-  - api/:         HTTP orchestration (ProtocolHandler, server)
-  - protocols/:   Response builders (OpenAI, Anthropic)
-  - transport/:   SSE transport utilities
-  - engine.py:    Facade (InferenceEngine), Value Object (GenerationRequest)
-  - sample.py:    Strategy pattern (TemperatureStrategy, TopKStrategy, TopPStrategy, FrequencyPenaltyStrategy)
+Subpackages:
+  - cache/:     KV cache (buffers, strategies, pool)
+  - runtime/:   Execution + sampling (executor, CUDA graph, sampling strategies)
+  - task/:      Request lifecycle + performance metrics
+  - network/:   HTTP protocol handling (server, protocol, OpenAI/Anthropic builders)
+
+Modules:
+  - scheduler.py:  Continuous batching loop
+  - workspace.py:  Pre-allocated GPU buffers
+  - engine.py:     Facade (InferenceEngine)
 """
 
-from astrai.inference.api import (
+from astrai.inference.cache import (
+    Allocator,
+    KVCache,
+    KVStorage,
+    PagePool,
+    RadixCache,
+    ReqToTokenPool,
+    TaskCacheManager,
+    page_hash,
+)
+from astrai.inference.engine import InferenceEngine
+from astrai.inference.network import (
     AnthropicMessage,
     BaseToolParser,
     ChatCompletionRequest,
@@ -25,25 +39,10 @@ from astrai.inference.api import (
     get_app,
     run_server,
 )
-from astrai.inference.api.anthropic import AnthropicResponseBuilder
-from astrai.inference.api.openai import OpenAIResponseBuilder
-from astrai.inference.core import (
-    STOP,
-    Allocator,
-    Executor,
-    InferenceScheduler,
-    KVCache,
-    KVStorage,
-    PagePool,
-    PrefixCache,
-    ReqToTokenPool,
-    Task,
-    TaskManager,
-    TaskStatus,
-    page_hash,
-)
-from astrai.inference.engine import GenerationRequest, InferenceEngine
-from astrai.inference.sample import (
+from astrai.inference.network.anthropic import AnthropicResponseBuilder
+from astrai.inference.network.openai import OpenAIResponseBuilder
+from astrai.inference.runtime.executor import Executor
+from astrai.inference.runtime.sample import (
     BaseSamplingStrategy,
     FrequencyPenaltyStrategy,
     SamplingPipeline,
@@ -52,10 +51,11 @@ from astrai.inference.sample import (
     TopPStrategy,
     sample,
 )
+from astrai.inference.scheduler import InferenceScheduler
+from astrai.inference.task import STOP, Task, TaskManager, TaskStatus
 
 __all__ = [
     "InferenceEngine",
-    "GenerationRequest",
     "InferenceScheduler",
     "Executor",
     "STOP",
@@ -66,8 +66,9 @@ __all__ = [
     "KVCache",
     "KVStorage",
     "PagePool",
-    "PrefixCache",
+    "RadixCache",
     "ReqToTokenPool",
+    "TaskCacheManager",
     "page_hash",
     "sample",
     "BaseSamplingStrategy",

astrai/inference/cache/__init__.py

@@ -0,0 +1,27 @@
+"""KV cache subsystem: buffers, strategies, pool management."""
+
+from astrai.inference.cache.buffer import KVCache, KVStorage, ReqToTokenPool
+from astrai.inference.cache.pool import PagePool, TaskCacheManager, page_hash
+from astrai.inference.cache.strategy import (
+    AllocationStrategy,
+    Allocator,
+    ContiguousStrategy,
+    PagedStrategy,
+    RadixCache,
+    TaskCacheState,
+)
+
+__all__ = [
+    "KVCache",
+    "KVStorage",
+    "ReqToTokenPool",
+    "Allocator",
+    "RadixCache",
+    "TaskCacheState",
+    "AllocationStrategy",
+    "ContiguousStrategy",
+    "PagedStrategy",
+    "PagePool",
+    "TaskCacheManager",
+    "page_hash",
+]

astrai/inference/cache/buffer.py

@@ -0,0 +1,106 @@
+"""Physical KV cache buffers.
+
+Layer 1 — ``KVStorage``:   flat token-level K/V GPU buffers [n_layers, size, n_kv_heads, head_dim]
+Layer 2 — ``ReqToTokenPool``: index table [req_idx, pos] → physical token slot
+Layer 3 — ``KVCache``:       pure dataclass passed to the model for direct buffer access
+
+These classes have no knowledge of tasks, allocation policies, or scheduling.
+They are the "dumb" physical storage layer.
+"""
+
+import threading
+from dataclasses import dataclass
+from typing import List, Optional
+
+import torch
+from torch import Tensor
+
+
+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.int32, 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.
+    """
+
+    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
+    q_tile_to_batch: Optional[Tensor] = None
+    q_tile_to_index: Optional[Tensor] = None
+    decode_o_part: Optional[Tensor] = None
+    decode_ml_part: Optional[Tensor] = None
+    decode_out: Optional[Tensor] = None

astrai/inference/cache/pool.py

@@ -0,0 +1,382 @@
+"""KV cache orchestration: PagePool + TaskCacheManager.
+
+PagePool owns the physical buffers (``KVStorage`` + ``ReqToTokenPool``)
+and wires them to an allocation strategy.  It assembles the ``KVCache``
+dataclass passed to the model forward.
+
+TaskCacheManager owns the ``task_id`` → ``TaskCacheState`` mapping and
+delegates physical slot allocation to the strategy, and KV bind to the pool.
+
+See ``cache_buffer.py`` for the raw buffer primitives and ``cache_strategy.py``
+for the allocation policies.
+"""
+
+from dataclasses import dataclass
+from typing import Dict, List, Optional
+
+import torch
+
+from astrai.inference.cache.buffer import KVCache, KVStorage, ReqToTokenPool
+from astrai.inference.cache.strategy import (
+    AllocationStrategy,
+    Allocator,
+    ContiguousStrategy,
+    PagedStrategy,
+    RadixCache,
+    TaskCacheState,
+)
+from astrai.inference.workspace import Q_TILE_ROWS, InferenceWorkspace
+
+# Re-export everything so existing ``from astrai.inference.cache import ...``
+# continues to work unchanged after the file split.
+__all__ = [
+    "KVCache",
+    "KVStorage",
+    "ReqToTokenPool",
+    "Allocator",
+    "RadixCache",
+    "AllocationStrategy",
+    "ContiguousStrategy",
+    "PagedStrategy",
+    "PagePool",
+    "TaskCacheManager",
+    "TaskCacheState",
+    "page_hash",
+]
+
+# ---- helpers ----
+
+
+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 = parent_hash
+    for i in range(start, end):
+        h = (h * 31 + token_ids[i]) & 0xFFFFFFFFFFFFFFFF
+    return h
+
+
+def _is_steady_increment(
+    prev_sig: Optional[tuple],
+    prev_vals: Optional[List[int]],
+    cur_sig: tuple,
+    cur_vals: List[int],
+) -> bool:
+    return (
+        prev_sig is not None
+        and prev_vals is not None
+        and prev_sig == cur_sig
+        and len(prev_vals) == len(cur_vals)
+        and all(c == p + 1 for c, p in zip(cur_vals, prev_vals))
+    )
+
+
+# ---- task-scoped bind state ----
+@dataclass
+class _BindState:
+    """Cached bind metadata for steady-state decode increment detection."""
+
+    sig: tuple
+    seq_lens: List[int]
+
+
+# ---- pool + manager ----
+
+
+class PagePool:
+    """Physical KV cache: buffers + req-to-token table + allocation strategy + bind.
+
+    Does not know about tasks — task lifecycle is managed by
+    :class:`TaskCacheManager`, which holds a reference to this pool.
+    """
+
+    def __init__(
+        self,
+        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
+
+        self.contiguous = n_tokens is None
+        self.n_tokens = max_batch_size * max_seq_len if self.contiguous else n_tokens
+        if self.n_tokens > torch.iinfo(torch.int32).max:
+            raise ValueError("KV cache token count exceeds the int32 slot index limit")
+
+        self._storage = KVStorage(
+            self.n_tokens, n_layers, n_kv_heads, head_dim, device, dtype
+        )
+        self._req_pool = ReqToTokenPool(max_batch_size, max_seq_len, device)
+
+        if self.contiguous:
+            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,
+                    dtype=torch.int32,
+                    device=device,
+                )
+            self._strategy: AllocationStrategy = ContiguousStrategy()
+        else:
+            n_pages = self.n_tokens // page_size
+            alloc = Allocator(n_pages)
+            prefix = RadixCache(page_size) if page_size > 1 else None
+            if prefix is not None:
+                alloc.on_evict = prefix.evict
+            self._strategy = PagedStrategy(
+                alloc, prefix, page_size, self._req_pool, device
+            )
+
+    @property
+    def strategy(self) -> AllocationStrategy:
+        return self._strategy
+
+    @property
+    def req_pool(self) -> ReqToTokenPool:
+        return self._req_pool
+
+    def bind_tasks(
+        self,
+        req_indices: List[int],
+        seq_lens: List[int],
+        workspace: InferenceWorkspace,
+        device: Optional[torch.device] = None,
+        start_pos: Optional[int] = None,
+        incremental: bool = False,
+    ) -> KVCache:
+        """Assemble the ``KVCache`` metadata for a batch of tasks.
+
+        Args:
+            req_indices: request slot indices (from ``ReqToTokenPool``).
+            seq_lens:    current sequence length per task.
+            workspace:   pre-allocated fixed-shape buffers (CUDA-graph safe).
+            start_pos:   if set, produce **prefill** cache (full q_len range).
+                         If ``None``, produce **decode** cache (last position).
+            incremental: if ``True``, reuse workspace state from previous step
+                         by incrementing counters in-place (decode hot path).
+
+        Returns:
+            ``KVCache`` dataclass with the correct output shapes for the
+            attention backend (prefill: ``[B, q_len]``, decode: ``[B, 1]``).
+        """
+        if device is None:
+            device = workspace.device
+        b = len(req_indices)
+
+        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
+
+        if incremental:
+            sl_buf[:b] += 1
+            kvp_buf[: b + 1] += inc_buf[: b + 1]
+        else:
+            rpi_buf[:b].copy_(
+                torch.tensor(req_indices, dtype=torch.int32, 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]
+
+        if start_pos is not None:
+            # Packed prefill concatenates each request's query tokens.
+            q_lens = [seq_len - start_pos for seq_len in seq_lens]
+            if any(q_len <= 0 for q_len in q_lens):
+                raise ValueError("prefill sequence lengths must exceed start_pos")
+            out_cache_loc = torch.cat(
+                [
+                    self._req_pool.req_to_token[
+                        req_pool_indices[i], start_pos : seq_lens[i]
+                    ]
+                    for i in range(b)
+                ]
+            )
+            workspace.qo_indptr[: b + 1].zero_()
+            workspace.qo_indptr[1 : b + 1].copy_(
+                torch.tensor(q_lens, dtype=torch.int32, device=device).cumsum(0)
+            )
+            qo_indptr = workspace.qo_indptr[: b + 1]
+            tile_batches = []
+            tile_indices = []
+            for batch, q_len in enumerate(q_lens):
+                n_tiles = (q_len + Q_TILE_ROWS - 1) // Q_TILE_ROWS
+                tile_batches.extend([batch] * n_tiles)
+                tile_indices.extend(range(n_tiles))
+            n_tiles = len(tile_batches)
+            workspace.q_tile_to_batch[:n_tiles].copy_(
+                torch.tensor(tile_batches, dtype=torch.int32, device=device)
+            )
+            workspace.q_tile_to_index[:n_tiles].copy_(
+                torch.tensor(tile_indices, dtype=torch.int32, device=device)
+            )
+            q_tile_to_batch = workspace.q_tile_to_batch[:n_tiles]
+            q_tile_to_index = workspace.q_tile_to_index[:n_tiles]
+            decode_o_part = decode_ml_part = decode_out = None
+        else:
+            # ---- decode: out_cache_loc is a single column (last position) ----
+            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].reshape(-1)
+            workspace.qo_indptr[: b + 1].copy_(inc_buf[: b + 1])
+            qo_indptr = workspace.qo_indptr[: b + 1]
+            q_tile_to_batch = q_tile_to_index = 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,
+            q_tile_to_batch=q_tile_to_batch,
+            q_tile_to_index=q_tile_to_index,
+            decode_o_part=decode_o_part,
+            decode_ml_part=decode_ml_part,
+            decode_out=decode_out,
+        )
+
+
+class TaskCacheManager:
+    """Task ↔ KV slot lifecycle manager.
+
+    Sole owner of ``task_id → TaskCacheState``.  Delegates physical slot
+    allocation to the strategy (via ``pool.strategy``) and KV bind to
+    ``pool.bind_tasks()``.
+
+    Usage::
+
+        pool = PagePool(...)
+        mgr = TaskCacheManager(pool)
+        mgr.task_alloc("req_1", [101, 202, 303])
+        ...
+        kv = mgr.bind(["req_1"], workspace)
+    """
+
+    def __init__(self, pool: PagePool):
+        self._pool = pool
+        self._strategy = pool.strategy
+        self._req_pool = pool.req_pool
+        self._max_seq_len = pool.max_seq_len
+        self._states: Dict[str, TaskCacheState] = {}
+        self._bind_state: Optional[_BindState] = None
+        self._bind_was_steady = False
+
+    # -- public task lifecycle --
+
+    def task_alloc(self, task_id: str, prompt_ids: List[int]) -> bool:
+        self._bind_state = None
+        req_slots = self._req_pool.alloc(1)
+        if req_slots is None:
+            return False
+        state = TaskCacheState(req_idx=req_slots[0])
+        self._states[task_id] = state
+        if not self._strategy.alloc(state, prompt_ids):
+            self._rollback(state, task_id)
+            return False
+        self._strategy.write_indices(state, prompt_ids)
+        state.length = len(prompt_ids)
+        return True
+
+    def task_free(self, task_id: str):
+        self._bind_state = None
+        state = self._states.pop(task_id, None)
+        if state is None:
+            return
+        self._strategy.free(state)
+        self._req_pool.free([state.req_idx])
+
+    def task_extend(self, task_id: str, pos: int) -> bool:
+        state = self._states.get(task_id)
+        if state is None or pos >= self._max_seq_len:
+            return False
+        if not self._strategy.extend(state, pos):
+            return False
+        state.length = pos + 1
+        return True
+
+    def task_cached(self, task_id: str) -> int:
+        state = self._states.get(task_id)
+        return state.cached if state is not None else 0
+
+    def task_record_hashes(
+        self, task_id: str, prompt_ids: List[int], start_logical_page: int = 0
+    ):
+        state = self._states.get(task_id)
+        if state is not None:
+            self._strategy.record_hashes(state, prompt_ids, start_logical_page)
+
+    @staticmethod
+    def task_cacheable_ids(task_id: str, prompt_ids: List[int], output_ids: List[int]):
+        return list(prompt_ids) + list(output_ids[:-1])
+
+    # -- bind (assemble KVCache for the model forward) --
+
+    def bind(
+        self,
+        task_ids: List[str],
+        workspace: InferenceWorkspace,
+        device: Optional[torch.device] = None,
+        start_pos: Optional[int] = None,
+    ) -> KVCache:
+        """Build ``KVCache`` for an ordered list of task IDs."""
+        states = [self._states[tid] for tid in task_ids]
+        req_indices = [s.req_idx for s in states]
+        seq_lens = [s.length for s in states]
+        sig = tuple(req_indices)
+
+        prev = self._bind_state
+        incremental = (
+            start_pos is None
+            and prev is not None
+            and _is_steady_increment(prev.sig, prev.seq_lens, sig, seq_lens)
+        )
+        self._bind_state = _BindState(sig, list(seq_lens))
+        self._bind_was_steady = incremental
+
+        return self._pool.bind_tasks(
+            req_indices,
+            seq_lens,
+            workspace,
+            device=device,
+            start_pos=start_pos,
+            incremental=incremental,
+        )
+
+    @property
+    def bind_was_steady(self) -> bool:
+        return self._bind_was_steady
+
+    # -- internals --
+
+    def _rollback(self, state: TaskCacheState, task_id: str):
+        self._strategy.free(state)
+        self._req_pool.free([state.req_idx])
+        self._states.pop(task_id, None)

astrai/inference/cache/strategy.py

@@ -0,0 +1,320 @@
+"""KV cache allocation layer.
+
+Encapsulates the physical slot allocation policy, isolated from GPU buffers
+and task lifecycle management.
+
+- ``TaskCacheState``: data contract between strategy and manager (per-task slot state)
+- ``Allocator``:       bitmask-based page allocator with LRU eviction
+- ``RadixCache``:      page-granular prefix index (exact token match)
+- ``AllocationStrategy``: ABC for physical slot allocation
+- ``ContiguousStrategy``: statically partitioned, no dynamic allocation
+- ``PagedStrategy``:    dynamic paged allocation from a shared pool
+"""
+
+import threading
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Callable, Dict, List, Optional, OrderedDict
+
+import torch
+
+from astrai.inference.cache.buffer import ReqToTokenPool
+
+# ---- data contract: per-task slot state ----
+
+
+@dataclass
+class TaskCacheState:
+    """Per-task cache allocation state.
+
+    Co-locates all task-owned cache metadata so the alloc/free/extend
+    lifecycle is atomic.  Owned by ``TaskCacheManager``, consumed by
+    every ``AllocationStrategy`` method.
+    """
+
+    req_idx: int
+    length: int = 0
+    cached: int = 0
+    pages: List[int] = field(default_factory=list)
+
+
+# ---- allocation primitives ----
+
+
+class Allocator:
+    """Bitmask-based page allocator with ref-counting and LRU eviction."""
+
+    def __init__(self, n_pages: int):
+        self._free_mask = (1 << n_pages) - 1
+        self._refs: List[int] = [0] * n_pages
+        self._lru: OrderedDict[int, None] = OrderedDict()
+        self.on_evict: Optional[Callable[[int], None]] = None
+        self._lock = threading.Lock()
+
+    def alloc(self) -> int:
+        with self._lock:
+            if self._free_mask:
+                lsb = self._free_mask & -self._free_mask
+                idx = lsb.bit_length() - 1
+                self._free_mask ^= lsb
+                self._refs[idx] = 1
+                return idx
+            if self._lru:
+                idx, _ = self._lru.popitem(last=False)
+                if self.on_evict:
+                    self.on_evict(idx)
+                self._refs[idx] = 1
+                self._free_mask &= ~(1 << idx)
+                return idx
+            return -1
+
+    def free(self, idx: int, keep_cached: bool = False):
+        with self._lock:
+            self._refs[idx] -= 1
+            if self._refs[idx] == 0:
+                if keep_cached:
+                    self._lru[idx] = None
+                else:
+                    self._free_mask |= 1 << idx
+
+    def inc_ref(self, idx: int):
+        with self._lock:
+            self._refs[idx] += 1
+            self._lru.pop(idx, None)
+
+    def ref_count(self, idx: int) -> int:
+        with self._lock:
+            return self._refs[idx]
+
+    def touch(self, idx: int):
+        with self._lock:
+            if idx in self._lru:
+                self._lru.move_to_end(idx)
+
+
+class RadixNode:
+    """A page-aligned edge in the CPU-side prefix radix trie."""
+
+    __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] = {}
+        self._lock = threading.Lock()
+
+    def evict(self, idx: int):
+        with self._lock:
+            node = self._page_to_node.pop(idx, 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_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):
+                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 or child.page_idx is None:
+                    break
+                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:
+            full_pages = len(token_ids) // self._page_size
+            if logical_page_idx >= full_pages:
+                return
+            old = self._page_to_node.pop(page_idx, None)
+            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)
+            node.page_idx = page_idx
+            self._page_to_node[page_idx] = node
+
+    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 AllocationStrategy(ABC):
+    """Physical slot allocation policy.
+
+    Subclasses implement the actual allocation semantics.  This ABC declares
+    the contract; there are no default implementations.
+    """
+
+    @abstractmethod
+    def alloc(self, state: TaskCacheState, prompt_ids: List[int]) -> bool: ...
+
+    @abstractmethod
+    def free(self, state: TaskCacheState) -> None: ...
+
+    @abstractmethod
+    def extend(self, state: TaskCacheState, pos: int) -> bool: ...
+
+    @abstractmethod
+    def write_indices(self, state: TaskCacheState, prompt_ids: List[int]) -> None: ...
+
+    @abstractmethod
+    def record_hashes(
+        self,
+        state: TaskCacheState,
+        prompt_ids: List[int],
+        start: int,
+    ) -> None: ...
+
+
+class ContiguousStrategy(AllocationStrategy):
+    """Static contiguous allocation: slots are pre-assigned at pool init.
+
+    No dynamic allocation or prefix caching.  All operations are no-ops
+    because ``ReqToTokenPool`` is pre-filled with contiguous ranges.
+    """
+
+    def alloc(self, state: TaskCacheState, prompt_ids: List[int]) -> bool:
+        return True
+
+    def free(self, state: TaskCacheState) -> None:
+        pass
+
+    def extend(self, state: TaskCacheState, pos: int) -> bool:
+        return True
+
+    def write_indices(self, state: TaskCacheState, prompt_ids: List[int]) -> None:
+        pass
+
+    def record_hashes(
+        self,
+        state: TaskCacheState,
+        prompt_ids: List[int],
+        start: int,
+    ) -> None:
+        pass
+
+
+class PagedStrategy(AllocationStrategy):
+    """Dynamic paged allocation from a shared bitmask pool.
+
+    ``page_size`` is a parameter, not a separate strategy: at ``page_size=1``
+    each allocated page *is* one token slot (``page * 1 + 0``), and prefix
+    caching is simply disabled (``prefix=None``).  The unified page formula
+    ``pages[page_idx] * page_size + offset`` holds for both.
+    """
+
+    def __init__(
+        self,
+        alloc: Allocator,
+        prefix: Optional[RadixCache],
+        page_size: int,
+        req_pool: ReqToTokenPool,
+        device,
+    ):
+        self._alloc = alloc
+        self._prefix = prefix
+        self._page_size = page_size
+        self._req_pool = req_pool
+        self._device = device
+
+    def alloc(self, state: TaskCacheState, prompt_ids: List[int]) -> bool:
+        if self._prefix is not None:
+            hits = self._prefix.lookup(prompt_ids)
+            state.cached = len(hits) * self._page_size
+            for p in hits:
+                self._alloc.inc_ref(p)
+            state.pages = list(hits)
+
+        remaining = len(prompt_ids) - state.cached
+        if remaining <= 0:
+            return True
+        n_new = (remaining + self._page_size - 1) // self._page_size
+        for _ in range(n_new):
+            p = self._alloc.alloc()
+            if p < 0:
+                return False
+            state.pages.append(p)
+        return True
+
+    def free(self, state: TaskCacheState) -> None:
+        if self._prefix is not None:
+            for p in state.pages:
+                keep = self._prefix.has_page(p)
+                self._alloc.free(p, keep_cached=keep)
+                if not keep:
+                    self._prefix.evict(p)
+        else:
+            for p in state.pages:
+                self._alloc.free(p)
+
+    def extend(self, state: TaskCacheState, pos: int) -> bool:
+        page_idx = pos // self._page_size
+        if page_idx >= len(state.pages):
+            p = self._alloc.alloc()
+            if p < 0:
+                return False
+            state.pages.append(p)
+        offset = pos % self._page_size
+        self._req_pool.req_to_token[state.req_idx, pos] = (
+            state.pages[page_idx] * self._page_size + offset
+        )
+        return True
+
+    def write_indices(self, state: TaskCacheState, prompt_ids: List[int]) -> None:
+        total = len(prompt_ids)
+        for pos in range(total):
+            page_idx = pos // self._page_size
+            offset = pos % self._page_size
+            if page_idx < len(state.pages):
+                self._req_pool.req_to_token[state.req_idx, pos] = (
+                    state.pages[page_idx] * self._page_size + offset
+                )
+
+    def record_hashes(
+        self,
+        state: TaskCacheState,
+        prompt_ids: List[int],
+        start: int,
+    ) -> None:
+        if self._prefix is None:
+            return
+        full = len(prompt_ids) // self._page_size
+        for i in range(start, min(full, len(state.pages))):
+            self._prefix.record(state.pages[i], prompt_ids, i)

astrai/inference/core/__init__.py

@@ -1,30 +0,0 @@
-"""Inference core: cache, executor, scheduler, task management."""
-
-from astrai.inference.core.cache import (
-    Allocator,
-    KVCache,
-    KVStorage,
-    PagePool,
-    PrefixCache,
-    ReqToTokenPool,
-    page_hash,
-)
-from astrai.inference.core.executor import Executor
-from astrai.inference.core.scheduler import InferenceScheduler
-from astrai.inference.core.task import STOP, Task, TaskManager, TaskStatus
-
-__all__ = [
-    "Allocator",
-    "KVCache",
-    "KVStorage",
-    "PagePool",
-    "PrefixCache",
-    "ReqToTokenPool",
-    "page_hash",
-    "Executor",
-    "InferenceScheduler",
-    "STOP",
-    "Task",
-    "TaskManager",
-    "TaskStatus",
-]

astrai/inference/core/cache.py

@@ -1,501 +0,0 @@
-"""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 PrefixCache (content 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 dataclasses import dataclass
-from typing import Callable, Dict, List, Optional
-
-import torch
-from torch import Tensor
-
-
-def page_hash(token_ids: List[int], page_idx: int, page_size: int) -> int:
-    start = page_idx * page_size
-    end = min(start + page_size, len(token_ids))
-    h = 0
-    for i in range(start, end):
-        h = (h * 31 + token_ids[i]) & 0xFFFFFFFFFFFFFFFF
-    return h
-
-
-class Allocator:
-    """Bitmask-based page allocator with ref-counting and LRU eviction."""
-
-    def __init__(self, n_pages: int):
-        self._free_mask = (1 << n_pages) - 1
-        self._refs: List[int] = [0] * n_pages
-        self._lru: OrderedDict[int, None] = OrderedDict()
-        self.on_evict: Optional[Callable[[int], None]] = None
-        self._lock = threading.Lock()
-
-    def alloc(self) -> int:
-        with self._lock:
-            if self._free_mask:
-                lsb = self._free_mask & -self._free_mask
-                idx = lsb.bit_length() - 1
-                self._free_mask ^= lsb
-                self._refs[idx] = 1
-                return idx
-            if self._lru:
-                idx, _ = self._lru.popitem(last=False)
-                if self.on_evict:
-                    self.on_evict(idx)
-                self._refs[idx] = 1
-                self._free_mask &= ~(1 << idx)
-                return idx
-            return -1
-
-    def free(self, idx: int, keep_cached: bool = False):
-        with self._lock:
-            self._refs[idx] -= 1
-            if self._refs[idx] == 0:
-                if keep_cached:
-                    self._lru[idx] = None
-                else:
-                    self._free_mask |= 1 << idx
-
-    def inc_ref(self, idx: int):
-        with self._lock:
-            self._refs[idx] += 1
-            self._lru.pop(idx, None)
-
-    def ref_count(self, idx: int) -> int:
-        with self._lock:
-            return self._refs[idx]
-
-    def touch(self, idx: int):
-        with self._lock:
-            if idx in self._lru:
-                self._lru.move_to_end(idx)
-
-
-class PrefixCache:
-    """Hash-based prefix matching: maps page hashes to physical page indices."""
-
-    def __init__(self, page_size: int):
-        self._page_size = page_size
-        self._page_to_hash: Dict[int, int] = {}
-        self._hash_to_page: Dict[int, int] = {}
-        self._lock = threading.Lock()
-
-    def evict(self, idx: int):
-        with self._lock:
-            h = self._page_to_hash.pop(idx, None)
-            if h is not None:
-                self._hash_to_page.pop(h, None)
-
-    def has_page(self, idx: int) -> bool:
-        with self._lock:
-            return idx in self._page_to_hash
-
-    def lookup(self, token_ids: List[int]) -> List[int]:
-        with self._lock:
-            full_pages = len(token_ids) // self._page_size
-            hits: List[int] = []
-            for i in range(full_pages):
-                h = page_hash(token_ids, i, self._page_size)
-                p = self._hash_to_page.get(h)
-                if p is None:
-                    break
-                hits.append(p)
-            return hits
-
-    def record(self, page_idx: int, token_ids: List[int], logical_page_idx: int):
-        with self._lock:
-            h = page_hash(token_ids, logical_page_idx, self._page_size)
-            old_h = self._page_to_hash.pop(page_idx, None)
-            if old_h is not None:
-                self._hash_to_page.pop(old_h, None)
-            self._page_to_hash[page_idx] = h
-            self._hash_to_page[h] = page_idx
-
-
-class 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
-        page_table: [batch, max_len] — precomputed gather indices for decode;
-            None for prefill or when not yet computed.
-        decode_mask: [batch, max_len] bool — precomputed position validity
-            mask for decode; None for prefill or single-batch decode.
-    """
-
-    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
-    page_table: Optional[Tensor] = None
-    decode_mask: Optional[Tensor] = None
-
-
-class PagePool:
-    """Top-level KV cache manager.
-
-    Combines KVStorage + ReqToTokenPool + Allocator + PrefixCache.
-
-    Args:
-        n_layers: Number of transformer layers.
-        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).
-    """
-
-    def __init__(
-        self,
-        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
-
-        self.contiguous = n_tokens is None
-        if self.contiguous:
-            self.n_tokens = max_batch_size * max_seq_len
-        else:
-            self.n_tokens = n_tokens
-
-        self._storage = KVStorage(
-            self.n_tokens, n_layers, n_kv_heads, head_dim, device, dtype
-        )
-        self._req_pool = ReqToTokenPool(max_batch_size, max_seq_len, device)
-
-        if self.contiguous:
-            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[PrefixCache] = None
-        else:
-            n_pages = self.n_tokens // page_size
-            self._alloc = Allocator(n_pages)
-            self._prefix = PrefixCache(page_size) if page_size > 1 else None
-            if self._prefix is not None:
-                self._alloc.on_evict = self._prefix.evict
-
-        self._task_req: Dict[str, int] = {}
-        self._task_len: Dict[int, int] = {}
-        self._task_cached: Dict[str, int] = {}
-        self._task_slots: Dict[str, List[int]] = {}
-        self._task_pages: Dict[str, List[int]] = {}
-        self._lock = threading.Lock()
-
-    # ---- task lifecycle ----
-
-    def task_alloc(self, task_id: str, prompt_ids: List[int]) -> bool:
-        req_slots = self._req_pool.alloc(1)
-        if req_slots is None:
-            return False
-        req_idx = req_slots[0]
-        self._task_req[task_id] = req_idx
-
-        if self.contiguous:
-            self._task_len[req_idx] = len(prompt_ids)
-            self._task_cached[task_id] = 0
-            return True
-
-        n_tokens_needed = len(prompt_ids)
-        cached = 0
-
-        if self._prefix is not None:
-            hits = self._prefix.lookup(prompt_ids)
-            cached = len(hits) * self.page_size
-            for p in hits:
-                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
-                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._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):
-        req_idx = self._task_req.pop(task_id, None)
-        if req_idx is None:
-            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:
-        req_idx = self._task_req.get(task_id)
-        if req_idx is None:
-            return False
-
-        if self.contiguous:
-            return pos < self.max_seq_len
-
-        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]
-        else:
-            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
-                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
-
-        self._task_len[req_idx] = pos + 1
-        return True
-
-    def task_cached(self, task_id: str) -> int:
-        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
-    ):
-        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, min(full_pages, len(pages))):
-            self._prefix.record(pages[i], prompt_ids, i)
-
-    # ---- bind for forward ----
-
-    def bind_tasks(
-        self,
-        task_ids: List[str],
-        seq_lens: List[int],
-        device: torch.device,
-        start_pos: Optional[int] = None,
-    ) -> KVCache:
-        req_indices = [self._task_req[tid] for tid in task_ids]
-        req_pool_indices = torch.tensor(req_indices, dtype=torch.long, device=device)
-        seq_lens_t = torch.tensor(seq_lens, dtype=torch.long, device=device)
-
-        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
-            ]
-            page_table = None
-            decode_mask = None
-        else:
-            write_pos = seq_lens_t - 1
-            out_cache_loc = self._req_pool.req_to_token[
-                req_pool_indices, write_pos
-            ].unsqueeze(-1)
-            ml = max(seq_lens)
-            page_table = self._req_pool.req_to_token[req_pool_indices, :ml]
-            if len(task_ids) > 1:
-                decode_mask = (
-                    torch.arange(ml, device=device)[None, :] < seq_lens_t[:, None]
-                )
-            else:
-                decode_mask = 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),
-            page_table=page_table,
-            decode_mask=decode_mask,
-        )
-
-    # ---- 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

astrai/inference/core/executor.py

@@ -1,174 +0,0 @@
-import logging
-from typing import List, Optional
-
-import torch
-
-from astrai.inference.core.cache import PagePool
-from astrai.inference.core.task import Task
-from astrai.inference.sample import sample
-from astrai.model.automodel import AutoModel
-from astrai.tokenize.tokenizer import AutoTokenizer
-
-logger = logging.getLogger(__name__)
-
-
-class Executor:
-    """Model forward passes for prefill and decode phases."""
-
-    def __init__(
-        self,
-        model: AutoModel,
-        tokenizer: AutoTokenizer,
-        kv_cache: PagePool,
-        device: Optional[str] = None,
-        dtype: Optional[torch.dtype] = None,
-    ):
-        self.model = model
-        self.tokenizer = tokenizer
-        self.kv_cache = kv_cache
-        self.device = device or next(model.parameters()).device
-        self.dtype = dtype or next(model.parameters()).dtype
-
-    def execute_prefill(self, tasks: List[Task], prompt_len: int, start_pos: int = 0):
-        if start_pos >= prompt_len:
-            return
-
-        tasks = sorted(tasks, key=lambda t: t.task_id)
-        batch_sz = len(tasks)
-
-        input_ids = torch.tensor(
-            [t.prompt_ids[start_pos:prompt_len] for t in tasks],
-            dtype=torch.long,
-            device=self.device,
-        )
-
-        task_ids = [t.task_id for t in tasks]
-        position_ids = (
-            torch.arange(start_pos, prompt_len, dtype=torch.long, device=self.device)
-            .unsqueeze(0)
-            .expand(batch_sz, -1)
-        )
-        input_mask = position_ids.unsqueeze(-1) >= torch.arange(
-            prompt_len, device=self.device
-        )
-
-        with torch.inference_mode():
-            self.model(
-                input_ids,
-                input_mask=input_mask,
-                position_ids=position_ids,
-                kv_cache=self.kv_cache.bind_tasks(
-                    task_ids, [prompt_len] * batch_sz, self.device, start_pos=start_pos
-                ),
-            )
-
-    def execute_decode(
-        self, tasks: List[Task], return_logprobs: bool = False
-    ) -> List[int]:
-        """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 []
-
-        input_ids = torch.tensor(
-            [t.output_ids[-1] if t.output_ids else t.prompt_ids[-1] for t in tasks],
-            dtype=torch.long,
-            device=self.device,
-        )
-
-        position_ids = torch.tensor(
-            [t.next_pos for t in tasks], dtype=torch.long, device=self.device
-        )
-        total_len = max(t.next_pos for t in tasks) + 1
-        input_mask = position_ids[:, None, None] >= torch.arange(
-            total_len, device=self.device
-        )
-
-        task_ids = [t.task_id for t in tasks]
-
-        temperatures = torch.tensor([t.temperature for t in tasks], device=self.device)
-        top_ks = torch.tensor([t.top_k for t in tasks], device=self.device)
-        top_ps = torch.tensor([t.top_p for t in tasks], device=self.device)
-        freq_penalties = torch.tensor(
-            [t.frequency_penalty for t in tasks], device=self.device
-        )
-
-        has_freq = bool((freq_penalties != 0).any())
-        if has_freq:
-            history_lists = []
-            history_lens = []
-            for t in tasks:
-                window = t.rep_window
-                prompt_part = t.prompt_ids[-window:]
-                ids = prompt_part + t.output_ids
-                history_lists.append(ids)
-                history_lens.append(len(ids))
-
-            max_len = max(history_lens) if history_lens else 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, h in enumerate(history_lists):
-                L = history_lens[i]
-                padded_ids[i, :L] = torch.as_tensor(
-                    h, dtype=torch.long, device=self.device
-                )
-                padded_mask[i, :L] = True
-        else:
-            padded_ids = None
-            padded_mask = None
-
-        with torch.inference_mode():
-            outputs = self.model(
-                input_ids.unsqueeze(1),
-                input_mask=input_mask,
-                kv_cache=self.kv_cache.bind_tasks(
-                    task_ids,
-                    [t.next_pos + 1 for t in tasks],
-                    self.device,
-                ),
-                position_ids=position_ids.unsqueeze(1),
-            )
-            logits = outputs["logits"][:, -1, :]
-
-        if return_logprobs:
-            tokens, logprobs = sample(
-                logits,
-                temperature=temperatures,
-                top_k=top_ks,
-                top_p=top_ps,
-                frequency_penalty=freq_penalties,
-                input_ids=padded_ids,
-                input_mask=padded_mask,
-                return_logprobs=True,
-            )
-            tokens_list = tokens.tolist()
-            logprobs_list = logprobs.tolist()
-            for t, lp in zip(tasks, logprobs_list):
-                t.output_logprobs.append(float(lp))
-            return list(zip(tokens_list, logprobs_list))
-
-        return sample(
-            logits,
-            temperature=temperatures,
-            top_k=top_ks,
-            top_p=top_ps,
-            frequency_penalty=freq_penalties,
-            input_ids=padded_ids,
-            input_mask=padded_mask,
-        ).tolist()

astrai/inference/core/scheduler.py

@@ -1,311 +0,0 @@
-import logging
-import threading
-import uuid
-from typing import Any, Dict, List, Optional, Tuple
-
-import torch
-
-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
-from astrai.tokenize.tokenizer import AutoTokenizer
-
-logger = logging.getLogger(__name__)
-
-
-class InferenceScheduler:
-    """Continuous batching loop: cleanup -> refill -> prefill -> decode (all groups)."""
-
-    def __init__(
-        self,
-        model: AutoModel,
-        tokenizer: AutoTokenizer,
-        max_batch_size: int = 16,
-        max_seq_len: Optional[int] = None,
-        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_position_embeddings is not None:
-            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_position_embeddings)"
-            )
-        self.device = device or next(model.parameters()).device
-        self.dtype = dtype or next(model.parameters()).dtype
-
-        head_dim = config.hidden_size // config.num_attention_heads
-
-        if cache is not None:
-            self._cache = cache
-        else:
-            self._cache = PagePool(
-                n_layers=config.num_hidden_layers,
-                n_kv_heads=config.num_key_value_heads,
-                head_dim=head_dim,
-                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,
-        )
-
-        self._executor = Executor(
-            model=model,
-            tokenizer=tokenizer,
-            kv_cache=self._cache,
-            device=self.device,
-            dtype=self.dtype,
-        )
-
-        self._stop_event = threading.Event()
-        self._loop_thread: Optional[threading.Thread] = None
-
-    def add_task(self, prompt: str, **kwargs) -> str:
-        return self._task_mgr.add_task(prompt, **kwargs)
-
-    def remove_task(self, task_id: str):
-        for task in self._task_mgr.remove_task(task_id):
-            self._cache.task_free(task.task_id)
-
-    def get_stats(self) -> Dict[str, Any]:
-        return self._task_mgr.get_stats()
-
-    def _run_generation_loop(self):
-        stop_ids = self._task_mgr.tokenizer.stop_ids
-        cache = self._cache
-        try:
-            while not self._stop_event.is_set():
-                finished = self._task_mgr.remove_finished_tasks(stop_ids)
-                for task in finished:
-                    cache.task_free(task.task_id)
-
-                active = self._task_mgr.get_active_tasks()
-                available = self._task_mgr.max_batch_size - len(active)
-                if available > 0:
-                    candidates = self._task_mgr.pull_candidates(available)
-                    failed = []
-                    for task in candidates:
-                        if cache.task_alloc(task.task_id, task.prompt_ids):
-                            self._task_mgr.activate(task)
-                        else:
-                            failed.append(task)
-                    if failed:
-                        self._task_mgr.return_to_waiting(failed)
-
-                if not self._task_mgr.has_work():
-                    self._task_mgr.wait_for_tasks(timeout=1.0)
-                    continue
-
-                active = self._task_mgr.get_active_tasks()
-
-                to_prefill = [
-                    t
-                    for t in active
-                    if t.output_tokens == 0
-                    and cache.task_cached(t.task_id) < len(t.prompt_ids)
-                ]
-                if to_prefill:
-                    for t in to_prefill:
-                        t.input_tokens = len(t.prompt_ids)
-
-                    groups: Dict[Tuple[int, int], List[Task]] = {}
-                    for t in to_prefill:
-                        key = (
-                            len(t.prompt_ids),
-                            cache.task_cached(t.task_id),
-                        )
-                        groups.setdefault(key, []).append(t)
-
-                    for (prompt_len, start_pos), group in groups.items():
-                        self._executor.execute_prefill(group, prompt_len, start_pos)
-                        start_logical_page = start_pos // getattr(
-                            cache, "page_size", 64
-                        )
-                        for t in group:
-                            cache.task_record_hashes(
-                                t.task_id, t.prompt_ids, start_logical_page
-                            )
-
-                decode_tasks = active
-
-                valid: List[Task] = []
-                for t in decode_tasks:
-                    if cache.task_extend(t.task_id, t.next_pos):
-                        valid.append(t)
-                    else:
-                        t.status = TaskStatus.ABORTED
-                        self._task_mgr.invoke_callback(t.task_id, STOP)
-
-                if valid:
-                    next_tokens = self._executor.execute_decode(valid)
-
-                    for t, ntok in zip(valid, next_tokens):
-                        t.output_ids.append(ntok)
-                        t.output_tokens += 1
-                        new_text = t.decode_new_token(self._task_mgr.tokenizer)
-                        if new_text:
-                            self._task_mgr.invoke_callback(t.task_id, new_text)
-
-                    for t in valid:
-                        if t.is_finished(stop_ids):
-                            remaining = t.flush_remaining(self._task_mgr.tokenizer)
-                            if remaining:
-                                self._task_mgr.invoke_callback(t.task_id, remaining)
-                            self._task_mgr.invoke_callback(t.task_id, STOP)
-
-        except Exception as e:
-            self._stop_event.set()
-            logger.error(f"Scheduler loop crashed: {e}", exc_info=True)
-            for task in self._task_mgr.get_active_tasks():
-                self._task_mgr.invoke_callback(task.task_id, STOP)
-                cache.task_free(task.task_id)
-            for task in self._task_mgr.get_waiting_tasks():
-                self._task_mgr.invoke_callback(task.task_id, STOP)
-            self._task_mgr.clear_queues()
-
-    def start(self):
-        if self._loop_thread is not None and self._loop_thread.is_alive():
-            return
-        self._stop_event.clear()
-        t = threading.Thread(target=self._run_generation_loop, daemon=True)
-        t.start()
-        self._loop_thread = t
-
-    def stop(self):
-        self._stop_event.set()
-        self._task_mgr.wake()
-        if self._loop_thread is not None:
-            self._loop_thread.join(timeout=2.0)
-        self._loop_thread = None
-        for task in self._task_mgr.get_active_tasks():
-            self._task_mgr.invoke_callback(task.task_id, STOP)
-            self._cache.task_free(task.task_id)
-        for task in self._task_mgr.get_waiting_tasks():
-            self._task_mgr.invoke_callback(task.task_id, STOP)
-            self._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))
-            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]
-            prefill_groups: Dict[Tuple[int, int], List[Task]] = {}
-            for t in live:
-                key = (len(t.prompt_ids), cache.task_cached(t.task_id))
-                prefill_groups.setdefault(key, []).append(t)
-            for (prompt_len, start_pos), group in prefill_groups.items():
-                self._executor.execute_prefill(group, prompt_len, start_pos)
-
-            while live:
-                valid: List[Task] = []
-                for t in sorted(live, key=lambda x: x.task_id):
-                    if cache.task_extend(t.task_id, t.next_pos):
-                        valid.append(t)
-                    else:
-                        t.status = TaskStatus.ABORTED
-                if not valid:
-                    break
-
-                step_out = self._executor.execute_decode(
-                    valid, return_logprobs=return_logprobs
-                )
-                if return_logprobs:
-                    for t, (ntok, _lp) in zip(valid, step_out):
-                        t.output_ids.append(ntok)
-                        t.output_tokens += 1
-                else:
-                    for t, ntok in zip(valid, step_out):
-                        t.output_ids.append(ntok)
-                        t.output_tokens += 1
-
-                live = [t for t in valid 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

astrai/inference/engine.py

@@ -8,9 +8,10 @@ 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.extension import ATTN_BACKEND, AttentionBackend, get_backend
+from astrai.inference.cache import PagePool
+from astrai.inference.scheduler import InferenceScheduler
+from astrai.inference.task import STOP
 from astrai.tokenize import AutoTokenizer
 
 
@@ -64,44 +65,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,
-        frequency_penalty: float = 0.0,
-        rep_window: int = 64,
-        stream: bool = False,
-    ):
-        if not (isinstance(top_k, int) and top_k >= 0):
-            raise ValueError("top_k must be a non-negative integer")
-        if not (0.0 <= top_p <= 1.0):
-            raise ValueError("top_p must be a float between 0.0 and 1.0")
-        if not (isinstance(temperature, (int, float)) and temperature >= 0):
-            raise ValueError("temperature must be a non-negative number")
-        if not (
-            isinstance(frequency_penalty, (int, float))
-            and -2.0 <= frequency_penalty <= 2.0
-        ):
-            raise ValueError("frequency_penalty must be between -2.0 and 2.0")
-        if not (isinstance(rep_window, int) and rep_window > 0):
-            raise ValueError("rep_window must be a positive integer")
-
-        self.messages = messages
-        self.top_k = top_k
-        self.top_p = top_p
-        self.temperature = temperature
-        self.max_tokens = max_tokens
-        self.frequency_penalty = frequency_penalty
-        self.rep_window = rep_window
-        self.stream = stream
-
-
 class InferenceEngine:
     """Unified inference engine backed by continuous-batching scheduler."""
 
@@ -112,6 +75,8 @@ class InferenceEngine:
         max_batch_size: int = 1,
         max_seq_len: Optional[int] = None,
         cache: Optional[PagePool] = None,
+        enable_cuda_graph: bool = True,
+        backend: Optional[Union[str, ATTN_BACKEND, AttentionBackend, type]] = None,
     ):
         self.model = model
         self.tokenizer = tokenizer
@@ -121,6 +86,8 @@ class InferenceEngine:
             max_batch_size=max_batch_size,
             max_seq_len=max_seq_len,
             cache=cache,
+            enable_cuda_graph=enable_cuda_graph,
+            backend=backend,
         )
 
         self.scheduler.start()
@@ -146,28 +113,23 @@ class InferenceEngine:
         is_batch = isinstance(prompt, list)
         prompts = prompt if is_batch else [prompt]
 
-        if stream:
-            return self._generate_streaming(
-                prompts,
-                is_batch,
-                max_tokens,
-                temperature,
-                top_p,
-                top_k,
-                frequency_penalty,
-                rep_window,
-            )
-        else:
-            return self._generate_non_streaming(
-                prompts,
-                is_batch,
-                max_tokens,
-                temperature,
-                top_p,
-                top_k,
-                frequency_penalty,
-                rep_window,
-            )
+        if max_tokens is not None and max_tokens <= 0:
+            if stream:
+                return iter(())
+            results = [""] * len(prompts)
+            return results if is_batch else results[0]
+
+        return self._generate(
+            prompts,
+            is_batch,
+            stream,
+            max_tokens,
+            temperature,
+            top_p,
+            top_k,
+            frequency_penalty,
+            rep_window,
+        )
 
     def generate_async(
         self,
@@ -179,9 +141,10 @@ class InferenceEngine:
         frequency_penalty: float = 0.0,
         rep_window: int = 64,
     ) -> AsyncGenerator[str, None]:
-        sync_gen = self._generate_streaming(
+        sync_gen = self._generate(
             [prompt],
             False,
+            True,
             max_tokens,
             temperature,
             top_p,
@@ -193,51 +156,31 @@ class InferenceEngine:
         async def _agen():
             loop = asyncio.get_event_loop()
             while True:
-                token = await loop.run_in_executor(None, self._next_token, sync_gen)
-                if token is None:
+                try:
+                    token = await loop.run_in_executor(None, next, sync_gen)
+                except StopIteration:
                     break
                 yield token
 
         return _agen()
 
-    @staticmethod
-    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,
-            frequency_penalty=request.frequency_penalty,
-            rep_window=request.rep_window,
-        )
-
-    def _submit_tasks(
+    def _generate(
         self,
         prompts: List[str],
+        is_batch: bool,
+        stream: bool,
         max_tokens: Optional[int],
         temperature: float,
         top_p: float,
         top_k: int,
         frequency_penalty: float,
         rep_window: int,
-    ) -> Tuple[GenerateResult, List[str]]:
+    ) -> Union[Generator, str, List[str]]:
         n = len(prompts)
+        request_backend = get_backend(use_default=False)
         result = GenerateResult(count=n)
-        task_ids = []
-        for i, p in enumerate(prompts):
-            cb = self._make_callback(result, i)
-            task_id = self.scheduler.add_task(
+        task_ids = [
+            self.scheduler.add_task(
                 prompt=p,
                 max_tokens=max_tokens,
                 temperature=temperature,
@@ -245,99 +188,54 @@ class InferenceEngine:
                 top_k=top_k,
                 frequency_penalty=frequency_penalty,
                 rep_window=rep_window,
-                stream_callback=cb,
+                backend=request_backend,
+                stream_callback=lambda token, idx=i: result.append(token, idx),
             )
-            task_ids.append(task_id)
-        return result, task_ids
+            for i, p in enumerate(prompts)
+        ]
 
-    @staticmethod
-    def _make_callback(result: GenerateResult, idx: int):
-        def cb(token):
-            result.append(token, idx)
+        if not stream:
+            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]
 
-        return cb
-
-    def _generate_streaming(
-        self,
-        prompts: List[str],
-        is_batch: bool,
-        max_tokens: Optional[int],
-        temperature: float,
-        top_p: float,
-        top_k: int,
-        frequency_penalty: float,
-        rep_window: int,
-    ) -> Generator:
-        result, task_ids = self._submit_tasks(
-            prompts,
-            max_tokens,
-            temperature,
-            top_p,
-            top_k,
-            frequency_penalty,
-            rep_window,
-        )
-        n = len(prompts)
         remaining = n
         finished = [False] * n
 
         def gen():
             nonlocal remaining
-            try:
-                while remaining > 0:
-                    items = result.pop_all()
-                    for idx, token in items:
-                        if token is STOP:
-                            if not finished[idx]:
-                                finished[idx] = True
-                                remaining -= 1
-                        else:
-                            yield (idx, token) if is_batch else token
-                    if remaining > 0:
-                        result.wait(timeout=0.05)
-            finally:
-                for tid in task_ids:
-                    self.scheduler.remove_task(tid)
+            while remaining > 0:
+                items = result.pop_all()
+                for idx, token in items:
+                    if token is STOP:
+                        if not finished[idx]:
+                            finished[idx] = True
+                            remaining -= 1
+                    else:
+                        yield (idx, token) if is_batch else token
+                if remaining > 0:
+                    result.wait(timeout=0.05)
 
         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,
-        frequency_penalty: float,
-        rep_window: int,
-    ) -> Union[str, List[str]]:
-        result, task_ids = self._submit_tasks(
-            prompts,
-            max_tokens,
-            temperature,
-            top_p,
-            top_k,
-            frequency_penalty,
-            rep_window,
-        )
-
-        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()
 
+    @property
+    def backend_name(self) -> str:
+        return self.scheduler.backend_name
+
+    @property
+    def cuda_graph_enabled(self) -> bool:
+        return self.scheduler.cuda_graph_enabled
+
     def shutdown(self):
         self.scheduler.stop()
         if torch.cuda.is_available():

astrai/inference/metrics.py

@@ -0,0 +1,223 @@
+"""Unified per-task perf/stats: timing records, context-manager scopes, aggregate reporting."""
+
+import time
+from collections import deque
+from contextlib import contextmanager
+from dataclasses import dataclass
+from typing import Any, Deque, Dict, Generator, List, Literal, Optional
+
+
+@dataclass
+class TaskTiming:
+    """Timestamp snapshots and computed metrics for one generation task.
+
+    Created by :class:`MetricsCollector` at task-registration time;
+    updated via ``record`` / ``mark_finished``.
+    """
+
+    task_id: str
+    arrival_time: float
+    prefill_start_time: Optional[float] = None
+    first_token_time: Optional[float] = None
+    finish_time: Optional[float] = None
+    input_tokens: int = 0
+    output_tokens: int = 0
+    _decode_steps: int = 0
+    _decode_total_s: float = 0.0
+
+    # derived metrics
+
+    @property
+    def queue_wait_ms(self) -> Optional[float]:
+        if self.prefill_start_time is not None:
+            return (self.prefill_start_time - self.arrival_time) * 1000
+        return None
+
+    @property
+    def ttft_ms(self) -> Optional[float]:
+        if self.first_token_time is not None:
+            return (self.first_token_time - self.arrival_time) * 1000
+        return None
+
+    @property
+    def prefill_tps(self) -> Optional[float]:
+        if self.prefill_start_time is not None and self.first_token_time is not None:
+            d = self.first_token_time - self.prefill_start_time
+            if d > 0 and self.input_tokens > 0:
+                return self.input_tokens / d
+        return None
+
+    @property
+    def decode_tps(self) -> Optional[float]:
+        if self.first_token_time is not None and self.finish_time is not None:
+            d = self.finish_time - self.first_token_time
+            dt = self.output_tokens - 1
+            if dt > 0 and d > 0:
+                return dt / d
+        return None
+
+    @property
+    def decode_avg_ms(self) -> Optional[float]:
+        if self._decode_steps > 0 and self._decode_total_s > 0:
+            return (self._decode_total_s / self._decode_steps) * 1000
+        return None
+
+    @property
+    def e2e_latency_ms(self) -> Optional[float]:
+        if self.finish_time is not None:
+            return (self.finish_time - self.arrival_time) * 1000
+        return None
+
+    @property
+    def total_tps(self) -> Optional[float]:
+        if self.finish_time is not None:
+            total = self.input_tokens + self.output_tokens
+            d = self.finish_time - self.arrival_time
+            if total > 0 and d > 0:
+                return total / d
+        return None
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "task_id": self.task_id,
+            "input_tokens": self.input_tokens,
+            "output_tokens": self.output_tokens,
+            "queue_wait_ms": (
+                round(self.queue_wait_ms, 2) if self.queue_wait_ms is not None else None
+            ),
+            "ttft_ms": (round(self.ttft_ms, 2) if self.ttft_ms is not None else None),
+            "prefill_tps": (
+                round(self.prefill_tps, 2) if self.prefill_tps is not None else None
+            ),
+            "decode_tps": (
+                round(self.decode_tps, 2) if self.decode_tps is not None else None
+            ),
+            "decode_avg_ms": (
+                round(self.decode_avg_ms, 2) if self.decode_avg_ms is not None else None
+            ),
+            "total_tps": (
+                round(self.total_tps, 2) if self.total_tps is not None else None
+            ),
+            "e2e_latency_ms": (
+                round(self.e2e_latency_ms, 2)
+                if self.e2e_latency_ms is not None
+                else None
+            ),
+        }
+
+
+class MetricsCollector:
+    """Single-owner perf/stats hub for all generation tasks.
+
+    Usage::
+
+        metrics = MetricsCollector()
+        metrics.register(task_id, arrival_time)
+
+        with metrics.record(task_ids, "prefill"):
+            run_prefill(...)
+
+        metrics.mark_finished(task_id, input_tokens, output_tokens)
+
+        stats = metrics.get_stats()
+    """
+
+    def __init__(self, max_recent: int = 128):
+        self._timings: Dict[str, TaskTiming] = {}
+        self._completed: Deque[TaskTiming] = deque(maxlen=max_recent)
+
+        self._ttft_ms_sum = 0.0
+        self._ttft_ms_count = 0
+        self._decode_tps_sum = 0.0
+        self._decode_tps_count = 0
+        self._e2e_ms_sum = 0.0
+        self._e2e_ms_count = 0
+
+    def register(self, task_id: str):
+        """Create a timing record for a newly-created task."""
+        self._timings[task_id] = TaskTiming(task_id=task_id, arrival_time=time.time())
+
+    def mark_finished(self, task_id: str, input_tokens: int, output_tokens: int):
+        """Close timing for a finished/aborted task and move it to completed."""
+        timing = self._timings.pop(task_id, None)
+        if timing is None:
+            return
+        timing.finish_time = time.time()
+        timing.input_tokens = input_tokens
+        timing.output_tokens = output_tokens
+        self._completed.append(timing)
+        self._accumulate(timing)
+
+    def clear(self):
+        """Reset all state (e.g. on engine shutdown)."""
+        self._timings.clear()
+        self._completed.clear()
+        self._ttft_ms_sum = 0.0
+        self._ttft_ms_count = 0
+        self._decode_tps_sum = 0.0
+        self._decode_tps_count = 0
+        self._e2e_ms_sum = 0.0
+        self._e2e_ms_count = 0
+
+    # timing scopes
+
+    @contextmanager
+    def record(
+        self, task_ids: List[str], phase: Literal["prefill", "decode"]
+    ) -> Generator[None, None, None]:
+        tic = time.time()
+        yield
+        toc = time.time()
+        dt = toc - tic
+        for tid in task_ids:
+            t = self._timings.get(tid)
+            if t is None:
+                continue
+            if phase == "prefill":
+                t.prefill_start_time = tic
+                t.first_token_time = toc
+            elif phase == "decode":
+                t._decode_steps += 1
+                t._decode_total_s += dt
+
+    # access
+
+    def get_timing(self, task_id: str) -> Optional[TaskTiming]:
+        """Return the timing record for *task_id* (active or completed)."""
+        if task_id in self._timings:
+            return self._timings[task_id]
+        for t in self._completed:
+            if t.task_id == task_id:
+                return t
+        return None
+
+    # aggregate stats
+
+    def get_stats(self) -> Dict[str, Any]:
+        stats: Dict[str, Any] = {}
+        if self._ttft_ms_count > 0:
+            stats["avg_ttft_ms"] = round(self._ttft_ms_sum / self._ttft_ms_count, 2)
+        if self._decode_tps_count > 0:
+            stats["avg_decode_tps"] = round(
+                self._decode_tps_sum / self._decode_tps_count, 2
+            )
+        if self._e2e_ms_count > 0:
+            stats["avg_e2e_latency_ms"] = round(
+                self._e2e_ms_sum / self._e2e_ms_count, 2
+            )
+        if self._completed:
+            stats["recent_tasks"] = [t.to_dict() for t in self._completed]
+        return stats
+
+    # internal
+
+    def _accumulate(self, t: TaskTiming):
+        if t.ttft_ms is not None:
+            self._ttft_ms_sum += t.ttft_ms
+            self._ttft_ms_count += 1
+        if t.decode_tps is not None:
+            self._decode_tps_sum += t.decode_tps
+            self._decode_tps_count += 1
+        if t.e2e_latency_ms is not None:
+            self._e2e_ms_sum += t.e2e_latency_ms
+            self._e2e_ms_count += 1

astrai/inference/network/__init__.py

@@ -4,8 +4,7 @@
 lazy singleton FastAPI instance.
 """
 
-from astrai.inference.api.protocol import GenContext, ProtocolHandler, StopChecker
-from astrai.inference.api.server import (
+from astrai.inference.network.app import (
     AnthropicMessage,
     ChatCompletionRequest,
     ChatMessage,
@@ -15,7 +14,8 @@ from astrai.inference.api.server import (
     get_app,
     run_server,
 )
-from astrai.inference.api.tool_parser import (
+from astrai.inference.network.protocol import GenContext, ProtocolHandler, StopChecker
+from astrai.inference.network.tool_parser import (
     BaseToolParser,
     SimpleJsonToolParser,
     ToolParserFactory,

astrai/inference/network/anthropic.py

@@ -6,13 +6,13 @@ from typing import Any, Dict, List, Tuple, Union
 
 from pydantic import BaseModel
 
-from astrai.inference.api.protocol import (
+from astrai.inference.engine import InferenceEngine
+from astrai.inference.network.protocol import (
     GenContext,
     ResponseBuilder,
     StopInfo,
     sse_event,
 )
-from astrai.inference.engine import InferenceEngine
 
 
 def _extract_text(content: Union[str, List[Dict[str, Any]]]) -> str:

astrai/inference/network/app.py

@@ -18,10 +18,10 @@ import uvicorn
 from fastapi import APIRouter, FastAPI, HTTPException
 from pydantic import BaseModel, Field
 
-from astrai.inference.api.anthropic import AnthropicResponseBuilder
-from astrai.inference.api.openai import OpenAIResponseBuilder
-from astrai.inference.api.protocol import ProtocolHandler
 from astrai.inference.engine import InferenceEngine
+from astrai.inference.network.anthropic import AnthropicResponseBuilder
+from astrai.inference.network.openai import OpenAIResponseBuilder
+from astrai.inference.network.protocol import ProtocolHandler
 from astrai.model import AutoModel
 from astrai.tokenize import AutoTokenizer
 

astrai/inference/network/openai.py

@@ -7,14 +7,14 @@ from typing import Any, Dict, List, Optional, Tuple, Union
 
 from pydantic import BaseModel
 
-from astrai.inference.api.protocol import (
+from astrai.inference.engine import InferenceEngine
+from astrai.inference.network.protocol import (
     GenContext,
     ResponseBuilder,
     StopInfo,
     sse_event,
 )
-from astrai.inference.api.tool_parser import BaseToolParser, ToolParserFactory
-from astrai.inference.engine import InferenceEngine
+from astrai.inference.network.tool_parser import BaseToolParser, ToolParserFactory
 
 logger = logging.getLogger(__name__)
 

astrai/inference/network/protocol.py

@@ -181,12 +181,10 @@ class ProtocolHandler:
         self, agen: AsyncGenerator, ctx: GenContext, stop_sequences: List[str]
     ) -> Dict[str, Any]:
         checker = StopChecker(stop_sequences)
-        chunks: List[str] = []
         body = ""
         matched = None
 
         async for token in agen:
-            chunks.append(token)
             body += token
 
             matched = checker.check(body)
@@ -195,6 +193,5 @@ class ProtocolHandler:
 
             ctx.completion_tokens += 1
 
-        content = "".join(chunks)
         stop = StopInfo(matched=matched, body=body)
-        return self.builder.format_response(ctx, content, stop)
+        return self.builder.format_response(ctx, body, stop)

astrai/inference/runtime/__init__.py

@@ -0,0 +1,25 @@
+"""Execution primitives: forward passes, CUDA graphs, and sampling."""
+
+from astrai.inference.runtime.executor import Executor
+from astrai.inference.runtime.graph import CudaGraphContext
+from astrai.inference.runtime.sample import (
+    BaseSamplingStrategy,
+    FrequencyPenaltyStrategy,
+    SamplingPipeline,
+    TemperatureStrategy,
+    TopKStrategy,
+    TopPStrategy,
+    sample,
+)
+
+__all__ = [
+    "Executor",
+    "CudaGraphContext",
+    "BaseSamplingStrategy",
+    "FrequencyPenaltyStrategy",
+    "SamplingPipeline",
+    "TemperatureStrategy",
+    "TopKStrategy",
+    "TopPStrategy",
+    "sample",
+]

astrai/inference/runtime/executor.py

@@ -0,0 +1,421 @@
+import logging
+import time
+from contextlib import contextmanager
+from dataclasses import dataclass
+from typing import List, Optional
+
+import torch
+from torch import Tensor
+
+from astrai.extension.backend.attention import (
+    CudaBackend,
+    get_backend,
+)
+from astrai.inference.cache import PagePool, TaskCacheManager
+from astrai.inference.runtime.graph import CudaGraphContext
+from astrai.inference.runtime.sample import sample
+from astrai.inference.task import Task
+from astrai.inference.workspace import InferenceWorkspace
+from astrai.model.automodel import AutoModel
+
+logger = logging.getLogger(__name__)
+
+
+@contextmanager
+def timed(label: str, log: Optional[logging.Logger] = None):
+    """GPU-precise timer via CUDA events; falls back to perf_counter on CPU."""
+    log = log or logger
+    if not log.isEnabledFor(logging.DEBUG):
+        yield
+        return
+    use_cuda = torch.cuda.is_available()
+    if use_cuda:
+        start = torch.cuda.Event(enable_timing=True)
+        end = torch.cuda.Event(enable_timing=True)
+        start.record()
+    else:
+        tic = time.perf_counter()
+    yield
+    if use_cuda:
+        end.record()
+        torch.cuda.synchronize()
+        elapsed_ms = start.elapsed_time(end)
+    else:
+        elapsed_ms = (time.perf_counter() - tic) * 1000
+    log.debug("%s %.2fms", 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())
+
+
+@dataclass
+class DecodeSteadyState:
+    """Cached decode metadata for the steady-state case.
+
+    When the same ordered task set decodes one token per step, sampling
+    params and task signature are reused; only positions advance by 1.
+    """
+
+    task_sig: tuple
+    positions: list[int]
+    sampling_info: SamplingBatchInfo
+
+
+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,
+    task_cache: TaskCacheManager,
+    ws: InferenceWorkspace,
+    gctx: CudaGraphContext,
+    max_batch_size: int,
+    prompt_len: int = 1,
+    device: Optional[str] = None,
+):
+    dev = device or next(model.parameters()).device
+
+    # Prefill warmup: cuBLAS auto-tunes for the actual prompt-length tensor
+    # shapes on first call (F.linear is the dominant cost).  This also warms
+    # up the CUDA context (driver init) and compiles the graph-capture trace
+    # that follows.  Custom .so kernels do NOT need this — they are pre-built.
+    warmup_len = 64
+    tid = "_warmup_prefill"
+    if task_cache.task_alloc(tid, list(range(warmup_len))):
+        with (
+            torch.inference_mode(),
+            timed("warmup prefill", logger),
+        ):
+            kv = task_cache.bind([tid], ws, start_pos=0)
+            ids_in = torch.arange(warmup_len, device=dev)
+            pos_in = ids_in
+            model(
+                ids_in,
+                kv_cache=kv,
+                position_ids=pos_in,
+                fwd="prefill",
+            )
+        task_cache.task_free(tid)
+
+    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)
+
+    for b in batch_sizes:
+        task_ids = [f"_warmup_decode_{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 task_cache.task_alloc(tid, pt):
+                alloc_ok = False
+                break
+        if not alloc_ok:
+            for tid in task_ids:
+                task_cache.task_free(tid)
+            continue
+
+        with (
+            torch.inference_mode(),
+            timed(f"warmup decode b={b}", logger),
+        ):
+            for step in range(2):
+                seq_pos = step
+                ws.position_ids[:b] = seq_pos
+                for tid in task_ids:
+                    task_cache.task_extend(tid, seq_pos)
+                kv = task_cache.bind(task_ids, ws)
+                ids_buf = ws.fill_input_ids([step] * b)
+                gctx.forward(
+                    model,
+                    key=(b,),
+                    input_ids=ids_buf,
+                    kv_cache=kv,
+                    position_ids=ws.position_ids[:b],
+                    fwd="decode",
+                )
+
+        for tid in task_ids:
+            task_cache.task_free(tid)
+        torch.cuda.synchronize()
+
+
+class Executor:
+    """Model forward passes for prefill and decode phases."""
+
+    def __init__(
+        self,
+        model: AutoModel,
+        kv_cache: PagePool,
+        task_cache: TaskCacheManager,
+        device: Optional[str] = None,
+        dtype: Optional[torch.dtype] = None,
+        enable_cuda_graph: bool = True,
+    ):
+        self.model = model
+        self.kv_cache = kv_cache
+        self.task_cache = task_cache
+        self.device = device or next(model.parameters()).device
+        self.dtype = dtype or next(model.parameters()).dtype
+
+        # Per-step decode cache for the steady-state case (same ordered
+        # task set decodes one token per step).  Sampling params stay
+        # constant; only positions advance.
+        self._decode_cache: Optional[DecodeSteadyState] = 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
+        backend = get_backend()
+        self._graph_supported = backend.supports_graph() and 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()
+        if enable_cuda_graph:
+            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.task_cache,
+            self._workspace,
+            self._graph_ctx,
+            max_batch_size=self.kv_cache.max_batch_size,
+            device=self.device,
+        )
+
+    @property
+    def cuda_graph_enabled(self) -> bool:
+        return self._graph_ctx.enabled and self._graph_supported
+
+    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)
+
+        input_ids = torch.tensor(
+            [token for t in tasks for token in t.prompt_ids[start_pos:prompt_len]],
+            dtype=torch.long,
+            device=self.device,
+        )
+
+        task_ids = [t.task_id for t in tasks]
+        position_ids = torch.arange(
+            start_pos, prompt_len, dtype=torch.long, device=self.device
+        ).repeat(batch_sz)
+
+        with (
+            torch.inference_mode(),
+            timed(f"execute_prefill b={batch_sz} prompt_len={prompt_len}", logger),
+        ):
+            outputs = self.model(
+                input_ids,
+                position_ids=position_ids,
+                kv_cache=self.task_cache.bind(
+                    task_ids,
+                    self._workspace,
+                    start_pos=start_pos,
+                ),
+                fwd="prefill",
+            )
+            q_len = prompt_len - start_pos
+            logits = outputs["logits"][
+                torch.arange(1, batch_sz + 1, device=self.device) * q_len - 1
+            ]
+
+        return tasks, self._sample_logits(logits, tasks, return_logprobs)
+
+    def execute_decode(
+        self, tasks: List[Task], return_logprobs: bool = False
+    ) -> List[int]:
+        """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]
+
+        kv_cache = self.task_cache.bind(task_ids, ws)
+
+        task_sig = tuple(task_ids)
+        reuse_decode_state = (
+            self.task_cache.bind_was_steady
+            and self._decode_cache is not None
+            and self._decode_cache.task_sig == task_sig
+        )
+        if reuse_decode_state:
+            info = self._decode_cache.sampling_info
+            ws.position_ids[:b] += 1
+        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 = DecodeSteadyState(task_sig, cur_positions, info)
+
+        # ---- forward (graph replay or live run + capture) ----
+
+        use_graph = (
+            self._graph_ctx.enabled
+            and self._graph_supported
+            and get_backend().supports_graph()
+        )
+        key = (b,)
+        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,
+                    kv_cache=kv_cache,
+                    position_ids=ws.position_ids[:b],
+                    fwd="decode",
+                )
+            else:
+                outputs = self.model(
+                    input_ids,
+                    kv_cache=kv_cache,
+                    position_ids=ws.position_ids[:b],
+                    fwd="decode",
+                )
+            logits = outputs["logits"]
+
+        return self._sample_logits(logits, tasks, return_logprobs, info=info)

astrai/inference/runtime/graph.py

@@ -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

astrai/inference/runtime/sample.py

@@ -266,7 +266,7 @@ class SamplingPipeline(BaseSamplingStrategy):
     @staticmethod
     def _is_greedy(temperature: Union[float, Tensor]) -> bool:
         if isinstance(temperature, Tensor):
-            return temperature.numel() == 1 and temperature.item() == 0
+            return bool((temperature == 0).all())
         return temperature == 0
 
     @torch.inference_mode()
@@ -305,12 +305,12 @@ class SamplingPipeline(BaseSamplingStrategy):
             return tokens, chosen
 
         transformed = self.apply(logits, filter_value, input_ids, input_mask)
-        log_probs = torch.log_softmax(transformed.float(), dim=-1)
         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
 
@@ -363,24 +363,6 @@ def sample(
         ``True`` — a ``(token_ids, chosen_logprobs)`` tuple where
         ``chosen_logprobs`` has shape ``[batch]``.
     """
-    greedy = (
-        (
-            isinstance(temperature, Tensor)
-            and temperature.numel() == 1
-            and temperature.item() == 0
-        )
-        if isinstance(temperature, Tensor)
-        else temperature == 0
-    )
-
-    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)

astrai/inference/scheduler.py

@@ -0,0 +1,408 @@
+import logging
+import threading
+import uuid
+from contextlib import nullcontext
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import torch
+
+from astrai.extension import (
+    ATTN_BACKEND,
+    AttentionBackend,
+    attn_backend,
+    get_backend,
+)
+from astrai.inference.cache import PagePool, TaskCacheManager
+from astrai.inference.metrics import MetricsCollector
+from astrai.inference.runtime.executor import Executor
+from astrai.inference.task import STOP, Task, TaskManager, TaskStatus
+from astrai.model.automodel import AutoModel
+from astrai.tokenize.tokenizer import AutoTokenizer
+
+logger = logging.getLogger(__name__)
+
+
+class InferenceScheduler:
+    """Continuous batching loop: cleanup -> refill -> prefill -> decode (all groups)."""
+
+    def __init__(
+        self,
+        model: AutoModel,
+        tokenizer: AutoTokenizer,
+        max_batch_size: int = 16,
+        max_seq_len: Optional[int] = None,
+        device: Optional[str] = None,
+        dtype: Optional[torch.dtype] = None,
+        cache: Optional[PagePool] = None,
+        enable_cuda_graph: bool = True,
+        backend: Optional[Union[str, ATTN_BACKEND, AttentionBackend, type]] = None,
+    ):
+        config = model.config
+
+        if max_seq_len is not None:
+            self.max_seq_len = max_seq_len
+        elif config.max_position_embeddings is not None:
+            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_position_embeddings)"
+            )
+        self.device = device or next(model.parameters()).device
+        self.dtype = dtype or next(model.parameters()).dtype
+
+        head_dim = config.hidden_size // config.num_attention_heads
+
+        if cache is not None:
+            self._cache = cache
+        else:
+            self._cache = PagePool(
+                n_layers=config.num_hidden_layers,
+                n_kv_heads=config.num_key_value_heads,
+                head_dim=head_dim,
+                max_batch_size=max_batch_size,
+                max_seq_len=self.max_seq_len,
+                device=self.device,
+                dtype=self.dtype,
+            )
+
+        self._metrics = MetricsCollector()
+
+        self._task_cache = TaskCacheManager(self._cache)
+
+        self._task_mgr = TaskManager(
+            tokenizer=tokenizer,
+            max_batch_size=max_batch_size,
+            max_seq_len=self.max_seq_len,
+            metrics=self._metrics,
+        )
+
+        if backend is None:
+            self._backend = None
+            default_backend = get_backend()
+            self._backend_name = type(default_backend).__name__
+            with attn_backend(default_backend):
+                self._executor = Executor(
+                    model=model,
+                    kv_cache=self._cache,
+                    task_cache=self._task_cache,
+                    device=self.device,
+                    dtype=self.dtype,
+                    enable_cuda_graph=enable_cuda_graph,
+                )
+        else:
+            with attn_backend(backend):
+                self._backend = get_backend()
+                self._backend_name = type(self._backend).__name__
+                self._executor = Executor(
+                    model=model,
+                    kv_cache=self._cache,
+                    task_cache=self._task_cache,
+                    device=self.device,
+                    dtype=self.dtype,
+                    enable_cuda_graph=enable_cuda_graph,
+                )
+
+        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._task_cache.task_free(task.task_id)
+
+    def get_stats(self) -> Dict[str, Any]:
+        return self._task_mgr.get_stats()
+
+    @property
+    def backend_name(self) -> str:
+        return self._backend_name
+
+    @property
+    def cuda_graph_enabled(self) -> bool:
+        return self._executor.cuda_graph_enabled
+
+    def _backend_context(self):
+        if self._backend is None:
+            return nullcontext()
+        return attn_backend(self._backend)
+
+    @staticmethod
+    def _task_backend_groups(tasks: List[Task]):
+        groups = {}
+        for task in tasks:
+            groups.setdefault(task.backend, (task.backend, []))[1].append(task)
+        return groups.values()
+
+    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``.
+        """
+        to_prefill = [t for t in tasks if not t.prefill_done 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, Optional[AttentionBackend]], List[Task]] = {}
+            for t in to_prefill:
+                start_pos = min(
+                    self._task_cache.task_cached(t.task_id), len(t.prompt_ids) - 1
+                )
+                groups.setdefault((len(t.prompt_ids), start_pos, t.backend), []).append(
+                    t
+                )
+
+            for (prompt_len, start_pos, _), group in groups.items():
+                backend = group[0].backend
+                backend_context = (
+                    attn_backend(backend) if backend is not None else nullcontext()
+                )
+                with (
+                    backend_context,
+                    self._metrics.record([t.task_id for t in group], "prefill"),
+                ):
+                    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
+                    t.mark_prefill_done()
+                    prefilled_ids.add(t.task_id)
+                    produced.append(t)
+
+                start_logical_page = start_pos // self._cache.page_size
+                for t in group:
+                    self._task_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 self._task_cache.task_extend(t.task_id, t.next_pos):
+                decoded.append(t)
+            else:
+                t.status = TaskStatus.ABORTED
+                aborted.append(t)
+
+        for backend, group in self._task_backend_groups(decoded):
+            backend_context = (
+                attn_backend(backend) if backend is not None else nullcontext()
+            )
+            with (
+                backend_context,
+                self._metrics.record([t.task_id for t in group], "decode"),
+            ):
+                step_out = self._executor.execute_decode(
+                    group, return_logprobs=return_logprobs
+                )
+            for t, out in zip(group, step_out):
+                t.output_ids.append(out[0] if return_logprobs else out)
+                t.output_tokens += 1
+                t.advance_kv()
+                produced.append(t)
+
+        return produced, aborted
+
+    def _run_generation_loop(self):
+        stop_ids = self._task_mgr.tokenizer.stop_ids
+        try:
+            with self._backend_context():
+                while not self._stop_event.is_set():
+                    finished = self._task_mgr.remove_finished_tasks(stop_ids)
+                    for task in finished:
+                        if task.status == TaskStatus.FINISHED:
+                            self._task_cache.task_record_hashes(
+                                task.task_id,
+                                self._task_cache.task_cacheable_ids(
+                                    task.task_id, task.prompt_ids, task.output_ids
+                                ),
+                            )
+                        self._task_cache.task_free(task.task_id)
+
+                    active = self._task_mgr.get_active_tasks()
+                    available = self._task_mgr.max_batch_size - len(active)
+                    if available > 0:
+                        candidates = self._task_mgr.pull_candidates(available)
+                        failed = []
+                        for task in candidates:
+                            if self._task_cache.task_alloc(
+                                task.task_id, task.prompt_ids
+                            ):
+                                self._task_mgr.activate(task)
+                            else:
+                                failed.append(task)
+                        if failed:
+                            self._task_mgr.return_to_waiting(failed)
+
+                    if not self._task_mgr.has_work():
+                        self._task_mgr.wait_for_tasks(timeout=1.0)
+                        continue
+
+                    active = self._task_mgr.get_active_tasks()
+
+                    decoded, aborted = self._step(active)
+
+                    for t in aborted:
+                        self._task_mgr.invoke_callback(t.task_id, STOP)
+
+                    for t in decoded:
+                        new_text = t.decode_new_token(self._task_mgr.tokenizer)
+                        if new_text:
+                            self._task_mgr.invoke_callback(t.task_id, new_text)
+                        if t.is_finished(stop_ids):
+                            self._task_mgr.invoke_callback(t.task_id, STOP)
+
+        except Exception as e:
+            self._stop_event.set()
+            logger.error(f"Scheduler loop crashed: {e}", exc_info=True)
+            for task in self._task_mgr.get_active_tasks():
+                self._task_mgr.invoke_callback(task.task_id, STOP)
+                self._task_cache.task_free(task.task_id)
+            for task in self._task_mgr.get_waiting_tasks():
+                self._task_mgr.invoke_callback(task.task_id, STOP)
+            self._task_mgr.clear_queues()
+
+    def start(self):
+        if self._loop_thread is not None and self._loop_thread.is_alive():
+            return
+        self._stop_event.clear()
+        t = threading.Thread(target=self._run_generation_loop, daemon=True)
+        t.start()
+        self._loop_thread = t
+
+    def stop(self):
+        self._stop_event.set()
+        self._task_mgr.wake()
+        if self._loop_thread is not None:
+            self._loop_thread.join(timeout=2.0)
+        self._loop_thread = None
+        for task in self._task_mgr.get_active_tasks():
+            self._task_mgr.invoke_callback(task.task_id, STOP)
+            self._task_cache.task_free(task.task_id)
+        for task in self._task_mgr.get_waiting_tasks():
+            self._task_mgr.invoke_callback(task.task_id, STOP)
+            self._task_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
+        seq_cap = self.max_seq_len
+        request_backend = get_backend(use_default=False)
+
+        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,
+                backend=request_backend,
+            )
+            if not self._task_cache.task_alloc(task.task_id, task.prompt_ids):
+                tasks.append(None)
+                continue
+            task.input_tokens = len(task.prompt_ids)
+            self._metrics.register(task.task_id)
+            tasks.append(task)
+
+        try:
+            live = [t for t in tasks if t is not None]
+
+            with self._backend_context():
+                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:
+                    self._metrics.mark_finished(
+                        t.task_id, t.input_tokens, t.output_tokens
+                    )
+                    self._task_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

astrai/inference/task.py

@@ -1,16 +1,17 @@
-import logging
 import threading
 import time
 import uuid
 from collections import deque
 from enum import Enum
-from typing import Any, Callable, Deque, Dict, List, Optional
+from typing import TYPE_CHECKING, Any, Callable, Deque, Dict, List, Optional
 
 from tokenizers.decoders import DecodeStream
 
+from astrai.inference.metrics import MetricsCollector
 from astrai.tokenize.tokenizer import AutoTokenizer
 
-logger = logging.getLogger(__name__)
+if TYPE_CHECKING:
+    from astrai.extension import AttentionBackend
 
 STOP = object()
 
@@ -64,6 +65,7 @@ class Task:
         top_k: int = 50,
         frequency_penalty: float = 0.0,
         rep_window: int = 64,
+        backend: Optional["AttentionBackend"] = None,
     ):
         self.task_id = task_id
         self.prompt_ids = prompt_ids
@@ -73,16 +75,25 @@ class Task:
         self.top_k = top_k
         self.frequency_penalty = frequency_penalty
         self.rep_window = rep_window
+        self.backend = backend
 
         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._kv_len: int = 0
         self._decoder: Optional[StreamDecoder] = None
 
+    def mark_prefill_done(self):
+        """Prompt KV is materialized by prefill; first output sampled but
+        not yet written to KV."""
+        self._kv_len = self.input_tokens
+
+    def advance_kv(self):
+        """One more position written to KV (after a decode forward)."""
+        self._kv_len += 1
+
     def decode_new_token(self, tokenizer: AutoTokenizer) -> str:
         """Decode the last appended output token, buffering incomplete
         multi-byte sequences across calls.
@@ -93,19 +104,15 @@ class Task:
             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)
+        """KV position where the next decode step will write."""
+        return self._kv_len
+
+    @property
+    def prefill_done(self) -> bool:
+        """True when all prompt KV entries are materialized."""
+        return self._kv_len >= self.input_tokens > 0
 
     def is_finished(self, stop_ids: List[int]) -> bool:
         if self.max_tokens is not None and self.output_tokens >= self.max_tokens:
@@ -123,6 +130,7 @@ class TaskManager:
         tokenizer: AutoTokenizer,
         max_batch_size: int = 16,
         max_seq_len: int = 8192,
+        metrics: Optional["MetricsCollector"] = None,
     ):
         self.tokenizer = tokenizer
         self.max_batch_size = max_batch_size
@@ -138,6 +146,8 @@ class TaskManager:
         self._total_tasks = 0
         self._total_tokens = 0
 
+        self._metrics = metrics
+
     def add_task(
         self,
         prompt: str,
@@ -147,6 +157,7 @@ class TaskManager:
         top_k: int = 50,
         frequency_penalty: float = 0.0,
         rep_window: int = 64,
+        backend: Optional["AttentionBackend"] = None,
         stream_callback: Optional[Callable[[str], None]] = None,
     ) -> str:
         task_id = f"task_{int(time.time())}_{uuid.uuid4().hex[:8]}"
@@ -154,11 +165,6 @@ class TaskManager:
         if len(prompt_ids) > self.max_seq_len:
             prompt_ids = prompt_ids[-self.max_seq_len :]
 
-        if len(prompt_ids) > self.max_seq_len:
-            if stream_callback:
-                stream_callback(STOP)
-            return task_id
-
         if max_tokens is None:
             max_tokens = self.max_seq_len - len(prompt_ids)
         else:
@@ -173,6 +179,7 @@ class TaskManager:
             top_k=top_k,
             frequency_penalty=frequency_penalty,
             rep_window=rep_window,
+            backend=backend,
         )
 
         with self._lock:
@@ -181,6 +188,9 @@ class TaskManager:
             if stream_callback:
                 self._callbacks[task_id] = stream_callback
 
+        if self._metrics is not None:
+            self._metrics.register(task_id)
+
         self._task_event.set()
         return task_id
 
@@ -200,26 +210,33 @@ class TaskManager:
             cb(token)
 
     def get_stats(self) -> Dict[str, Any]:
-        return {
+        stats: Dict[str, Any] = {
             "total_tasks": self._total_tasks,
             "total_tokens": self._total_tokens,
             "active_tasks": len(self.active_tasks),
             "waiting_queue": len(self.waiting_queue),
         }
+        if self._metrics is not None:
+            stats.update(self._metrics.get_stats())
+        return stats
 
     def remove_finished_tasks(self, stop_ids: List[int]) -> List[Task]:
         with self._lock:
             finished = []
             for task in self.active_tasks:
                 if task.status == TaskStatus.ABORTED:
-                    task.finish_time = time.time()
                     finished.append(task)
                 elif task.is_finished(stop_ids):
                     task.status = TaskStatus.FINISHED
-                    task.finish_time = time.time()
                     finished.append(task)
                     self._total_tokens += task.output_tokens
 
+            if self._metrics is not None:
+                for task in finished:
+                    self._metrics.mark_finished(
+                        task.task_id, task.input_tokens, task.output_tokens
+                    )
+
             self.active_tasks = [
                 t
                 for t in self.active_tasks

astrai/inference/workspace.py

@@ -0,0 +1,159 @@
+"""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
+Q_TILE_ROWS = 64
+
+
+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.int32, 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
+        )
+        max_q_tiles = max_batch_size * ((max_seq_len + Q_TILE_ROWS - 1) // Q_TILE_ROWS)
+        self.q_tile_to_batch = torch.empty(
+            (max_q_tiles,), dtype=torch.int32, device=device
+        )
+        self.q_tile_to_index = torch.empty(
+            (max_q_tiles,), 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.int32, 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

astrai/logging.py

@@ -0,0 +1,27 @@
+import logging
+import os
+
+
+def setup_logging(level: str = "INFO"):
+    """Attach a StreamHandler to the ``astrai`` logger (idempotent).
+
+    Call once per process at the top of CLI scripts.
+    Set ``ASTR_LOG_LEVEL`` env var to override the default level.
+
+    Level names: ``DEBUG``, ``INFO``, ``WARNING``, ``ERROR``, ``CRITICAL``.
+    ``DEBUG`` enables per-step prefill/decode timing logs
+    (:func:`astrai.inference.runtime.executor.timed`).
+    """
+    logger = logging.getLogger("astrai")
+    if logger.handlers:
+        return
+    level_name = os.environ.get("ASTR_LOG_LEVEL", level).upper()
+    logger.setLevel(getattr(logging, level_name, logging.INFO))
+    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)

astrai/model/__init__.py

@@ -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

astrai/model/components/__init__.py

@@ -1,9 +1,9 @@
-from astrai.extension.rotary_backend import apply_rotary_emb
+from astrai.extension.backend.rotary 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,
@@ -14,6 +14,7 @@ __all__ = [
     "Linear",
     "RMSNorm",
     "MLP",
+    "DeepSeekMoE",
     "Embedding",
     "GQA",
     "MLA",

astrai/model/components/attention.py

@@ -5,10 +5,9 @@ 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.extension.backend import apply_rotary_emb, attention
 from astrai.factory import BaseFactory
-from astrai.inference.core.cache import KVCache
+from astrai.inference.cache import KVCache
 from astrai.model.components.linear import Linear
 from astrai.model.components.norm import RMSNorm
 
@@ -56,9 +55,7 @@ class GQA(nn.Module):
             self.gate = Linear(dim, dim)
 
     def _split_heads(self, x: Tensor, n_heads) -> Tensor:
-        batch_size, seq_len, _ = x.shape
-        x = x.reshape(batch_size, seq_len, n_heads, self.head_dim)
-        return x
+        return x.reshape(*x.shape[:-1], n_heads, self.head_dim)
 
     def forward(
         self,
@@ -67,6 +64,7 @@ class GQA(nn.Module):
         attn_mask: Tensor = None,
         kv_cache: Optional[KVCache] = None,
         is_causal: bool = False,
+        fwd: Optional[str] = None,
     ) -> Tensor:
         q = self._split_heads(self.q_proj(x), self.n_heads)
         k = self._split_heads(self.k_proj(x), self.n_kv_heads)
@@ -76,7 +74,9 @@ class GQA(nn.Module):
         if self.use_qk_norm:
             q, k = self.q_norm(q), self.k_norm(k)
 
-        sdqa_out = attention(q, k, v, kv_cache, self.layer_id, attn_mask, is_causal)
+        sdqa_out = attention(
+            q, k, v, kv_cache, self.layer_id, attn_mask, is_causal, fwd
+        ).reshape(*x.shape[:-1], self.dim)
 
         if self.use_gated_attention:
             sdqa_out = sdqa_out * F.sigmoid(self.gate(x))
@@ -141,17 +141,16 @@ class MLA(nn.Module):
         attn_mask: Tensor = None,
         kv_cache: Optional[KVCache] = None,
         is_causal: bool = False,
+        fwd: Optional[str] = None,
     ) -> Tensor:
-        bsz, seq_len, _ = x.size()
-
         q = self.q_proj(x)
-        q = q.view(bsz, seq_len, self.n_heads, self.head_dim)
+        q = q.reshape(*x.shape[:-1], self.n_heads, self.head_dim)
 
         kv_compressed = self.kv_a_proj(x)
         kv_compressed = self.kv_norm(kv_compressed)
 
         kv = self.kv_b_proj(kv_compressed)
-        kv = kv.view(bsz, seq_len, self.n_kv_heads, -1)
+        kv = kv.reshape(*x.shape[:-1], self.n_kv_heads, -1)
 
         k_nope, k_rope, v = torch.split(
             kv, [self.qk_nope_head_dim, self.qk_rope_head_dim, self.head_dim], dim=-1
@@ -171,7 +170,9 @@ class MLA(nn.Module):
             q = self.q_norm(q)
             k = self.k_norm(k)
 
-        attn_out = attention(q, k, v, kv_cache, self.layer_id, attn_mask, is_causal)
+        attn_out = attention(
+            q, k, v, kv_cache, self.layer_id, attn_mask, is_causal, fwd
+        ).reshape(*x.shape[:-1], self.dim)
 
         if self.use_gated_attention:
             attn_out = attn_out * F.sigmoid(self.gate(x))

astrai/model/components/decoder_block.py

@@ -1,15 +1,21 @@
 from dataclasses import asdict
-from typing import Optional
+from typing import Optional, TypedDict
 
 import torch.nn as nn
 from torch import Tensor
 
-from astrai.inference.core.cache import KVCache
+from astrai.inference.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__(self, config, layer_id: int):
         super().__init__()
@@ -26,7 +32,20 @@ class DecoderBlock(nn.Module):
         self.attention = AttnFactory.create(config.attn_type, **cfg, layer_id=layer_id)
         self.input_norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
         self.post_attention_norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
-        self.mlp = FFNFactory.create(config.ffn_type, **cfg)
+        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,
@@ -35,15 +54,23 @@ class DecoderBlock(nn.Module):
         attention_mask: Optional[Tensor] = None,
         kv_cache: Optional[KVCache] = None,
         is_causal: bool = False,
-    ) -> Tensor:
+        fwd: Optional[str] = None,
+    ) -> DecoderOutput:
         attn_output = self.attention(
             self.input_norm(x),
             rotary_emb,
             attention_mask,
             kv_cache,
             is_causal,
+            fwd,
         )
         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"),
+        }

astrai/model/components/mlp.py

@@ -1,3 +1,5 @@
+from typing import Optional, TypedDict
+
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -11,6 +13,28 @@ class FFNFactory(BaseFactory[nn.Module]):
     pass
 
 
+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, down_init_std: float = 0.02):
@@ -19,10 +43,10 @@ class MLP(nn.Module):
         self.gate = Linear(dim, dim_ffn)
         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")
@@ -36,6 +60,9 @@ class DeepSeekMoE(nn.Module):
         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
@@ -43,6 +70,16 @@ 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
@@ -50,51 +87,92 @@ class DeepSeekMoE(nn.Module):
 
         self.shared_experts = nn.ModuleList(
             [
-                MLP(dim, dim_ffn, down_init_std=down_init_std)
+                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, down_init_std=down_init_std)
+                MLP(dim, expert_dim_ffn, down_init_std=down_init_std)
                 for _ in range(n_routed_experts)
             ]
         )
 
-    def forward(self, x: Tensor) -> Tensor:
-        bsz, seq_len, dim = x.shape
+    def forward(self, x: Tensor) -> FFNOutput:
+        include_aux_loss = self.training and torch.is_grad_enabled()
+        shape = x.shape
+        dim = shape[-1]
         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)
-        return out
+        out = (shared_out + routed_output["hidden_states"]).view(shape)
+        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_weights / topk_weights.sum(dim=-1, keepdim=True)
+        topk_weights, topk_indices = torch.topk(router_probs, K, dim=-1, sorted=False)
+        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):
-            expert_mask = topk_indices == expert_idx
-            token_idx, k_idx = expert_mask.nonzero(as_tuple=True)
-            if token_idx.numel() == 0:
+        start = 0
+        for expert_idx, end in enumerate(boundaries):
+            if end == start:
                 continue
-            expert_input = x[token_idx]
-            expert_output = self.routed_experts[expert_idx](expert_input)
-            weights = topk_weights[token_idx, k_idx].unsqueeze(-1)
-            output.index_add_(0, token_idx, expert_output * weights)
+            expert_output = self.routed_experts[expert_idx](flat_tokens[start:end])[
+                "hidden_states"
+            ]
+            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,
+        }

astrai/model/components/rope.py

@@ -65,9 +65,12 @@ class RotaryEmbedding(nn.Module):
             [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)
-            )
+            if x.ndim == 2:
+                position_ids = torch.arange(x.size(0), device=x.device)
+            else:
+                position_ids = (
+                    torch.arange(x.size(1), device=x.device)
+                    .unsqueeze(0)
+                    .expand(x.size(0), -1)
+                )
         return self.freqs_cis[position_ids].float()

astrai/model/encoder.py

@@ -70,7 +70,7 @@ class EmbeddingEncoder(AutoModel):
         attn_mask = process_attention_mask(input_mask)
 
         for layer in self.layers:
-            x = layer(x, rotary_emb, attn_mask)
+            x = layer(x, rotary_emb, attn_mask)["hidden_states"]
 
         hidden_states = self.norm(x)
 

astrai/model/transformer.py

@@ -5,7 +5,7 @@ import torch.nn as nn
 from torch import Tensor
 
 from astrai.config.model_config import AutoRegressiveLMConfig
-from astrai.inference.core.cache import KVCache
+from astrai.inference.cache import KVCache
 from astrai.model.automodel import AutoModel, ModelFactory
 from astrai.model.components.decoder_block import DecoderBlock
 from astrai.model.components.embedding import Embedding
@@ -105,18 +105,48 @@ class AutoRegressiveLM(AutoModel):
         input_mask: Optional[Tensor] = None,
         kv_cache: Optional[KVCache] = None,
         position_ids: Optional[Tensor] = None,
+        fwd: Optional[str] = None,
     ) -> Dict[str, Tensor]:
-        assert input_ids.ndim == 2
+        if fwd is None:
+            if input_ids.ndim != 2:
+                raise ValueError("training input_ids must be [batch, seq_len]")
+            if kv_cache is not None:
+                raise ValueError("training forward does not accept a KV cache")
+        elif fwd in ("prefill", "decode"):
+            if input_ids.ndim != 1:
+                raise ValueError("inference input_ids must be packed [tokens]")
+            if kv_cache is None:
+                raise ValueError("inference forward requires a KV cache")
+        else:
+            raise ValueError(f"unsupported forward mode: {fwd}")
 
         x = self.embed_tokens(input_ids)
         rotary_emb = self.rotary_embedding(x, position_ids)
         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, kv_cache, use_sdpa_causal_mask)
+            layer_output = layer(
+                x,
+                rotary_emb,
+                attn_mask,
+                kv_cache,
+                use_sdpa_causal_mask,
+                fwd,
+            )
+            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

astrai/preprocessing/builder.py

@@ -416,7 +416,11 @@ class MultiOutputMaskBuilder(BaseMaskBuilder):
             return None
 
         result: dict = {}
-        any_output = False
+        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", [])
@@ -428,7 +432,6 @@ class MultiOutputMaskBuilder(BaseMaskBuilder):
                 if ids is None:
                     continue
                 result[output_key] = ids
-                any_output = True
                 continue
 
             list_field = spec.get("list_field", False)
@@ -444,7 +447,6 @@ class MultiOutputMaskBuilder(BaseMaskBuilder):
                 result[output_key] = ids
                 if mask is not None:
                     result[mask_key] = mask
-                any_output = True
                 continue
 
             ids, mask = self.renderer.process_sections(
@@ -460,9 +462,7 @@ class MultiOutputMaskBuilder(BaseMaskBuilder):
             elif "mask_key" in spec:
                 result[mask_key] = mask
 
-            any_output = True
-
-        if not any_output:
+        if not required_outputs or not required_outputs.issubset(result):
             return None
 
         result["domain"] = _extract_domain(item, config.output.domain_key)
@@ -474,6 +474,11 @@ class MultiOutputMaskBuilder(BaseMaskBuilder):
             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:
@@ -506,7 +511,7 @@ class MultiOutputMaskBuilder(BaseMaskBuilder):
 
         return [
             ({**result, "domain": _extract_domain(item, config.output.domain_key)})
-            if result
+            if required_outputs and required_outputs.issubset(result)
             else None
             for item, result in zip(items, results)
         ]

astrai/trainer/metric_util.py

@@ -1,3 +1,4 @@
+import math
 from typing import Dict
 
 import torch
@@ -27,6 +28,8 @@ class GradSNRTracker:
 
     SNR = E[g]^2 / Var(g) = E[g]^2 / (E[g^2] - E[g]^2)
 
+    The reported value is the power ratio in decibels: ``10 * log10(SNR)``.
+
     The tracker accumulates per-parameter EMA moments across optimizer steps.
     Call ``update`` after backward (before ``optimizer.step``) and read
     ``snr`` to get the aggregate SNR across all parameters.
@@ -64,7 +67,8 @@ class GradSNRTracker:
             noise = (v - m.pow(2)).clamp(min=0).sum().item()
             total_signal += signal
             total_noise += noise
-        return total_signal / (total_noise + self.eps)
+        snr = total_signal / (total_noise + self.eps)
+        return 10.0 * math.log10(max(snr, self.eps))
 
 
 def ctx_get_loss(ctx):
@@ -88,3 +92,23 @@ def ctx_get_grad_snr(ctx):
     if tracker is None:
         return None
     return tracker.snr
+
+
+def ctx_get_moe_aux_loss(ctx):
+    return ctx.strategy._moe_metrics.get("aux_loss")
+
+
+def ctx_get_router_entropy(ctx):
+    return ctx.strategy._moe_metrics.get("router_entropy")
+
+
+def ctx_get_dead_expert_fraction(ctx):
+    return ctx.strategy._moe_metrics.get("dead_expert_fraction")
+
+
+def ctx_get_load_imbalance_mean(ctx):
+    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")

astrai/trainer/rollout.py

@@ -6,7 +6,7 @@ Provides:
 - :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`
+  :class:`~astrai.inference.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)``
@@ -20,7 +20,7 @@ from typing import Dict, List, Optional, Tuple
 import torch
 from torch import Tensor
 
-from astrai.inference.core.scheduler import InferenceScheduler
+from astrai.inference.scheduler import InferenceScheduler
 
 
 @dataclass(kw_only=True)
@@ -101,7 +101,7 @@ 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`,
+    :meth:`~astrai.inference.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.

astrai/trainer/strategy.py

@@ -1,7 +1,7 @@
 """Training strategy implementations with factory pattern."""
 
 from abc import ABC, abstractmethod
-from typing import Callable, Dict, Union
+from typing import Callable, Dict, List, Optional, TypedDict, Union
 
 import torch
 import torch.nn as nn
@@ -9,10 +9,22 @@ import torch.nn.functional as F
 from torch import Tensor
 
 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
 
 
+class LossOutput(TypedDict):
+    loss: Tensor
+    metrics: Dict[str, float]
+
+
+class LogprobsOutput(TypedDict):
+    logprobs: Tensor
+    aux_loss: Optional[Tensor]
+    router_stats: Optional[List[RouterStats]]
+
+
 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()}
@@ -24,7 +36,7 @@ def get_logprobs(
     attn_mask: Tensor,
     loss_mask: Tensor,
     reduction: str,
-) -> Tensor:
+) -> LogprobsOutput:
     """Compute token-wise log probabilities from model outputs.
 
     Args:
@@ -46,10 +58,11 @@ def get_logprobs(
     shifted_input_ids = input_ids[:, 1:]
     shifted_loss_mask = loss_mask[:, 1:]
 
-    logits = model(
+    outputs = model(
         input_ids[:, :-1],
         attn_mask[:, :, :-1, :-1] if attn_mask.dim() == 4 else attn_mask[:, :-1],
-    )["logits"]
+    )
+    logits = outputs["logits"]
     log_probs = torch.log_softmax(logits.float(), dim=-1)
 
     token_logprobs = torch.gather(
@@ -57,13 +70,18 @@ def get_logprobs(
     ).squeeze(-1)
 
     if reduction == "mean":
-        return (token_logprobs * shifted_loss_mask).sum(dim=-1) / shifted_loss_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_loss_mask).sum(dim=-1)
+        logprobs = (token_logprobs * shifted_loss_mask).sum(dim=-1)
     else:
-        return token_logprobs * shifted_loss_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:
@@ -82,6 +100,68 @@ def make_doc_boundary_mask(position_ids: Tensor) -> Tensor:
     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.
 
@@ -102,6 +182,8 @@ class BaseStrategy(ABC):
         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
 
@@ -117,6 +199,35 @@ class BaseStrategy(ABC):
         """
         raise NotImplementedError
 
+    def compute_loss_output(self, batch: Dict[str, Tensor]) -> LossOutput:
+        return self._normalize_output(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.
 
@@ -148,22 +259,36 @@ class BaseStrategy(ABC):
         """
         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]) -> Tensor:
+    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(batch)
+            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(train_batch)
+        return self.compute_loss_output(train_batch)
 
 
 class StrategyFactory(BaseFactory["BaseStrategy"]):
@@ -190,6 +315,7 @@ 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__(
@@ -203,9 +329,13 @@ class SEQStrategy(BaseStrategy):
         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(),
@@ -213,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")
@@ -221,6 +356,7 @@ 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__(
@@ -234,6 +370,9 @@ class SFTStrategy(BaseStrategy):
         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, position_ids, loss_mask = (
             batch["input_ids"],
@@ -245,9 +384,10 @@ class SFTStrategy(BaseStrategy):
         ignore_index = -100
         input_mask = make_doc_boundary_mask(position_ids)
         target_ids = target_ids.masked_fill(~loss_mask, ignore_index)
-        logits = self.model(
+        outputs = self.model(
             input_ids=input_ids, position_ids=position_ids, input_mask=input_mask
-        )["logits"]
+        )
+        logits = outputs["logits"]
 
         loss = F.cross_entropy(
             input=logits.flatten(0, 1).float(),
@@ -256,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")
@@ -282,6 +427,9 @@ class DPOStrategy(BaseStrategy):
         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"]
@@ -297,22 +445,25 @@ class DPOStrategy(BaseStrategy):
         )[None, None, :, :]  # [1, 1, S, S]
         full_mask = key_pad & causal  # [B*2, 1, S, S] — composed
 
-        log_pi = get_logprobs(
+        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,
                 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] :]
@@ -325,7 +476,12 @@ 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
@@ -398,6 +554,9 @@ class GRPOStrategy(BaseStrategy):
             self.old_model.load_state_dict(state_dict)
 
     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)
         prompts = batch["prompts"]
         responses = batch["responses"]
@@ -438,16 +597,23 @@ class GRPOStrategy(BaseStrategy):
         # get_logprobs returns [B*G, S-1] (S = prompt_len + response_len).
         # Response token logprobs occupy the last ``response_len`` positions
         # (the first response token is predicted from the last prompt token).
-        token_log_probs_policy = get_logprobs(
+        policy_output = get_logprobs(
             self.model, full_sequences, attn_mask, full_masks, "none"
-        )[:, prompt_len - 1 :]
+        )
+        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():
-            token_log_probs_old = get_logprobs(
+            old_output = get_logprobs(
                 self.old_model, full_sequences, attn_mask, full_masks, "none"
-            )[:, prompt_len - 1 :]
-            token_log_probs_ref = get_logprobs(
+            )
+            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"
-            )[:, prompt_len - 1 :]
+            )
+            token_log_probs_ref = ref_output["logprobs"]
+            token_log_probs_ref = token_log_probs_ref[:, prompt_len - 1 :]
 
         # Reshape to [B, G, response_len]
         token_log_probs_policy = token_log_probs_policy.view(batch_size, group_size, -1)
@@ -480,9 +646,13 @@ class GRPOStrategy(BaseStrategy):
         kl_per_token = r - torch.log(r + eps) - 1.0
         kl_penalty = self.kl_coef * (kl_per_token * token_masks).sum() / token_count
 
-        total_loss = policy_loss + kl_penalty
-
-        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

astrai/trainer/train_callback.py

@@ -17,10 +17,15 @@ from astrai.parallel import only_on_rank
 from astrai.parallel.setup import get_current_device
 from astrai.serialization import Checkpoint
 from astrai.trainer.metric_util import (
+    ctx_get_dead_expert_fraction,
     ctx_get_grad_norm,
     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
@@ -257,14 +262,36 @@ class MetricCallback(TrainCallback):
             "val_loss": ctx_get_val_loss,
             "grad_norm": ctx_get_grad_norm,
             "grad_snr": ctx_get_grad_snr,
+            "moe_aux_loss": ctx_get_moe_aux_loss,
+            "router_entropy": ctx_get_router_entropy,
+            "dead_expert_fraction": ctx_get_dead_expert_fraction,
+            "load_imbalance_mean": ctx_get_load_imbalance_mean,
+            "load_imbalance_max": ctx_get_load_imbalance_max,
         }
 
     def _metrics(self, context: TrainContext, names):
-        return {
-            m: self._metric_funcs[m](context)
-            for m in names
-            if self._metric_funcs[m](context) is not None
-        }
+        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):
@@ -286,8 +313,8 @@ class MetricCallback(TrainCallback):
 
         with torch.no_grad():
             for batch in context.val_dataloader:
-                loss = context.strategy(batch)
-                total_loss += loss.item()
+                loss_output = context.strategy(batch)
+                total_loss += loss_output["loss"].item()
                 num_batches += 1
 
         if context.world_size > 1 and dist.is_initialized():

astrai/trainer/train_context.py

@@ -10,7 +10,7 @@ from torch.utils.data import DataLoader, random_split
 
 from astrai.config.train_config import TrainConfig
 from astrai.dataset import RDSampler
-from astrai.inference.core.scheduler import InferenceScheduler
+from astrai.inference.scheduler import InferenceScheduler
 from astrai.model.components.lora import inject_lora
 from astrai.parallel.executor import BaseExecutor, ExecutorFactory, create_ref_model
 from astrai.parallel.setup import get_current_device, get_rank, get_world_size
@@ -38,6 +38,7 @@ class TrainContext:
     epoch: int = field(default=0)
     consumed_samples: int = field(default=0)
     loss: float = field(default=0.0)
+    metrics: Dict[str, float] = field(default_factory=dict)
     grad_norm: Optional[float] = field(default=None)
     grad_snr_tracker: GradSNRTracker = field(default_factory=GradSNRTracker)
     val_dataloader: Optional[DataLoader] = field(default=None)
@@ -65,6 +66,15 @@ class TrainContext:
         )
 
 
+@dataclass
+class _PreloadedState:
+    model_config: dict = field(default_factory=dict)
+    state_dict: Optional[dict] = None
+    epoch: int = 0
+    consumed_samples: int = 0
+    checkpoint: Optional[Checkpoint] = None
+
+
 class TrainContextBuilder:
     def __init__(
         self,
@@ -80,212 +90,231 @@ class TrainContextBuilder:
         return self
 
     def build(self) -> TrainContext:
-        cfg = self.config
-        device = get_current_device()
+        # Resolve persisted state.
+        preloaded_state = self._load_preloaded_state()
 
-        executor = ExecutorFactory.create(
+        # Build the core training components and restore their persisted state.
+        executor = self._create_executor()
+        context = self._create_context(preloaded_state, executor)
+        self._prepare_model(context, executor, preloaded_state)
+        self._restore_optimizer_state(context)
+
+        # Resolve datasets.
+        train_dataset, val_dataset = self._get_datasets()
+        self._create_dataloaders(context, train_dataset, val_dataset)
+
+        # Strategies depend on the prepared model; online rollout depends on both.
+        strategy_kwargs = self._create_strategy(context, executor)
+        self._configure_rollout(context, strategy_kwargs)
+
+        return context
+
+    def _create_executor(self) -> BaseExecutor:
+        cfg = self.config
+        return ExecutorFactory.create(
             cfg.parallel_mode,
             grad_accum_steps=cfg.grad_accum_steps,
             **cfg.executor_kwargs,
         )
 
-        model_config = {}
+    def _load_preloaded_state(self) -> _PreloadedState:
+        cfg = self.config
+        state = _PreloadedState(
+            epoch=cfg.start_epoch,
+            consumed_samples=cfg.start_samples * get_world_size(),
+        )
         if self._param_path:
             config_path = Path(self._param_path) / "config.json"
             if config_path.exists():
-                model_config = load_json(config_path)
-
-        preloaded_state_dict = None
-        preloaded_epoch = cfg.start_epoch
-        preloaded_consumed = cfg.start_samples * get_world_size()
-        preloaded_checkpoint = None
-        if self._param_path:
+                state.model_config = load_json(config_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
+                state.state_dict = checkpoint.state_dict
+                state.model_config = checkpoint.config or state.model_config
                 if self._resume:
-                    preloaded_epoch = checkpoint.epoch
+                    state.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
+                    state.consumed_samples = (
+                        checkpoint.consumed_samples // per_step * per_step
+                    )
+                    state.checkpoint = checkpoint
+        if not state.model_config and hasattr(cfg.model_fn(), "config"):
+            state.model_config = cfg.model_fn().config.to_dict()
+        return state
 
-        if not model_config and hasattr(cfg.model_fn(), "config"):
-            model_config = cfg.model_fn().config.to_dict()
+    def _create_context(
+        self, state: _PreloadedState, executor: BaseExecutor
+    ) -> TrainContext:
+        return TrainContext(
+            world_size=get_world_size(),
+            rank=get_rank(),
+            config=self.config,
+            model_config=state.model_config,
+            executor=executor,
+            epoch=state.epoch,
+            consumed_samples=state.consumed_samples,
+            checkpoint=state.checkpoint,
+        )
 
-        def _before_wrap(m):
-            m = m.to(device=device)
+    def _prepare_model(
+        self, context: TrainContext, executor: BaseExecutor, state: _PreloadedState
+    ) -> None:
+        cfg = self.config
+        device = get_current_device()
+
+        def before_wrap(model):
+            model = model.to(device=device)
             if cfg.lora is not None:
                 inject_lora(
-                    m,
+                    model,
                     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
+            if state.state_dict is not None:
+                model.load_state_dict(state.state_dict, strict=False)
+            return model
 
-        def _after_wrap(m):
+        def after_wrap(model):
             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(
-            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,
-        )
+                model = torch.compile(model, mode=cfg.compile_mode)
+            return model
 
         context.model, context.optimizer, context.scheduler = executor.prepare(
             cfg.model_fn,
             cfg.optimizer_fn,
             cfg.scheduler_fn,
-            before_wrap=_before_wrap,
-            after_wrap=_after_wrap,
+            before_wrap=before_wrap,
+            after_wrap=after_wrap,
         )
 
-        train_dataset = cfg.dataset
-        val_dataset = cfg.val_dataset
+    def _get_datasets(self):
+        cfg = self.config
+        if cfg.val_dataset is not None or cfg.val_split is None:
+            return cfg.dataset, cfg.val_dataset
+        n_val = max(1, int(len(cfg.dataset) * cfg.val_split))
+        generator = torch.Generator().manual_seed(cfg.random_seed)
+        return random_split(
+            cfg.dataset, [len(cfg.dataset) - n_val, n_val], generator=generator
+        )
 
-        if val_dataset is None and cfg.val_split is not None:
-            n_total = len(cfg.dataset)
-            n_val = max(1, int(n_total * cfg.val_split))
-            n_train = n_total - n_val
-            generator = torch.Generator().manual_seed(cfg.random_seed)
-            train_dataset, val_dataset = random_split(
-                cfg.dataset, [n_train, n_val], generator=generator
+    def _create_dataloaders(
+        self, context: TrainContext, train_dataset, val_dataset
+    ) -> None:
+        cfg = self.config
+        sampler_offset = context.consumed_samples // context.world_size
+        if self._resume and sampler_offset > 0:
+            samples_per_replica = (
+                len(train_dataset) + context.world_size - 1
+            ) // context.world_size
+            if samples_per_replica > 0:
+                context.epoch = sampler_offset // samples_per_replica
+        context.dataloader = self._create_dataloader(
+            train_dataset, context.epoch, sampler_offset
+        )
+        if val_dataset is not None:
+            context.val_dataloader = self._create_dataloader(
+                val_dataset, 0, 0, shuffle=False
             )
 
-        sampler_offset = context.consumed_samples // context.world_size
-
-        if self._resume and sampler_offset > 0:
-            offset = context.world_size - 1
-            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
-
+    def _create_dataloader(
+        self, dataset, epoch: int, start_iter: int, shuffle: bool = True
+    ):
+        cfg = self.config
         sampler = RDSampler(
-            data_source=train_dataset,
-            start_epoch=context.epoch,
-            start_iter=sampler_offset,
+            dataset,
+            start_epoch=epoch,
+            start_iter=start_iter,
             seed=cfg.random_seed,
+            shuffle=shuffle,
         )
-        context.dataloader = DataLoader(
-            train_dataset,
+        loader_kwargs = dict(
+            dataset=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 val_dataset is not None:
-            val_sampler = RDSampler(
-                data_source=val_dataset,
-                start_epoch=0,
-                start_iter=0,
-                seed=cfg.random_seed,
-                shuffle=False,
-            )
-            context.val_dataloader = DataLoader(
-                val_dataset,
-                batch_size=cfg.batch_per_device,
-                sampler=val_sampler,
-                num_workers=cfg.num_workers,
-                pin_memory=cfg.pin_memory,
-                prefetch_factor=cfg.prefetch_factor,
-                collate_fn=cfg.collate_fn,
-            )
-
-        if context.checkpoint and context.checkpoint.extra:
-            extra = context.checkpoint.extra
-            for name in ("optimizer", "scheduler"):
-                if name in extra:
-                    obj = getattr(context, name, None)
-                    if obj is not None:
-                        obj.load_state_dict(extra[name])
-
-        strategy_kwargs = dict(cfg.extra_kwargs)
-
-        needs_ref = cfg.strategy in (
-            "dpo",
-            "grpo",
-            "online_grpo",
-            "online_dpo",
+        # PyTorch rejects prefetch_factor/persistent_workers when workers=0.
+        if cfg.num_workers > 0:
+            loader_kwargs["persistent_workers"] = cfg.persistent_workers
+            if cfg.prefetch_factor is not None:
+                loader_kwargs["prefetch_factor"] = cfg.prefetch_factor
+        return DataLoader(
+            **loader_kwargs,
         )
-        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
+    def _restore_optimizer_state(self, context: TrainContext) -> None:
+        if context.checkpoint and context.checkpoint.extra:
+            for name in ("optimizer", "scheduler"):
+                if (
+                    name in context.checkpoint.extra
+                    and getattr(context, name, None) is not None
+                ):
+                    getattr(context, name).load_state_dict(
+                        context.checkpoint.extra[name]
+                    )
+
+    def _create_strategy(self, context: TrainContext, executor: BaseExecutor) -> dict:
+        cfg = self.config
+        kwargs = dict(cfg.extra_kwargs)
+        kwargs.setdefault("moe_aux_loss_coef", cfg.moe_aux_loss_coef)
+        if cfg.strategy in ("dpo", "grpo", "online_grpo", "online_dpo"):
+            kwargs["ref_model"] = create_ref_model(
+                cfg.model_fn,
+                executor=executor,
+                model=context.model,
+                device=get_current_device(),
             )
-
-        if needs_old:
-            strategy_kwargs["old_model"] = create_ref_model(
-                cfg.model_fn, executor=executor, model=context.model, device=device
+        if cfg.strategy in ("grpo", "online_grpo"):
+            kwargs["old_model"] = create_ref_model(
+                cfg.model_fn,
+                executor=executor,
+                model=context.model,
+                device=get_current_device(),
             )
-
         context.strategy = StrategyFactory.create(
             cfg.strategy,
             model=context.model,
-            device=device,
+            device=get_current_device(),
             executor=executor,
-            **strategy_kwargs,
+            **kwargs,
         )
+        return 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,
+    def _configure_rollout(self, context: TrainContext, strategy_kwargs: dict) -> None:
+        cfg = self.config
+        if not cfg.strategy.startswith("online_"):
+            return
+        if not context.strategy.supports_online():
+            raise ValueError(
+                f"Strategy '{cfg.strategy}' does not support online rollout"
             )
-
-            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(
+        tokenizer = AutoTokenizer.from_pretrained(self._param_path)
+        group_size = strategy_kwargs.get("group_size", 1)
+        scheduler = InferenceScheduler(
+            model=context.model,
+            tokenizer=tokenizer,
+            max_batch_size=group_size * max(1, cfg.batch_per_device),
+            max_seq_len=getattr(context.model.config, "max_position_embeddings", None),
+        )
+        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,
+        )
+        context.strategy.set_rollout_runner(
+            RolloutRunner(
                 generator=generator,
-                reward_model=reward_model,
+                reward_model=cfg.reward_model_fn(),
                 rollout_interval=cfg.rollout_interval,
             )
-            context.strategy.set_rollout_runner(runner)
-
-        return context
+        )

astrai/trainer/trainer.py

@@ -82,9 +82,10 @@ class Trainer:
                         break
                     with executor.accumulate(context.model):
                         self._call_callbacks("on_batch_begin", context)
-                        loss = context.strategy(batch)
-                        context.loss = loss.item()
-                        stand_loss = loss / executor.grad_accum_steps
+                        loss_output = context.strategy(batch)
+                        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.consumed_samples += (
                             context.config.batch_per_device * context.world_size

csrc/CMakeLists.txt

@@ -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 fp8_mm)
+
+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()

csrc/build.py

@@ -1,75 +0,0 @@
-from pathlib import Path
-
-
-def cuda_toolkit_version() -> tuple[int, int] | None:
-    """Return ``(major, minor)`` of the nvcc on PATH, or ``None``.
-
-    Used by ``setup.py`` to detect nvcc/torch CUDA version mismatches
-    (e.g. nvcc 13.0 with a cu128 torch wheel) which cause cryptic ABI errors.
-    """
-    import shutil
-    import subprocess
-
-    nvcc = shutil.which("nvcc")
-    if nvcc is None:
-        return None
-    try:
-        out = subprocess.check_output(
-            [nvcc, "--version"], stderr=subprocess.STDOUT, text=True
-        )
-        for line in out.splitlines():
-            if "release" in line:
-                ver = line.split("release")[1].split(",")[0].strip()
-                major, minor = ver.split(".")
-                return (int(major), int(minor))
-    except Exception:
-        pass
-    return None
-
-
-def _arch_flags() -> list[str]:
-    import torch
-
-    if torch.cuda.is_available():
-        cap = torch.cuda.get_device_capability()
-    else:
-        cap = (8, 0)
-    ver = f"{cap[0]}{cap[1]}"
-    flags = [f"-gencode=arch=compute_{ver},code=sm_{ver}"]
-    # tensor-core mma path (mma.sync.m16n8k16.bf16) requires sm_80+; decide the
-    # kernel dispatch at build time via this define rather than at runtime.
-    if cap[0] < 8:
-        flags.append("-DASTRAI_NO_MMA")
-    return flags
-
-
-_kernels_dir = Path("csrc/kernels")
-REGISTRY: dict[str, dict] = {}
-
-CXX_FLAGS = ["-O3", "-funroll-loops"]
-NVCC_FLAGS = [
-    "-O3",
-    "--expt-relaxed-constexpr",
-    "--use_fast_math",
-    "--ptxas-options=-O3,-v",
-    "--extra-device-vectorization",
-    "--threads=16",
-]
-
-
-def register(name: str, sources: list[str] | None = None, **kwargs):
-    if sources is None:
-        sources = [str(_kernels_dir / f"{name}.cu")]
-    REGISTRY[name] = {
-        "sources": sources,
-        "cxx_flags": [*CXX_FLAGS],
-        "nvcc_flags": [*NVCC_FLAGS, *_arch_flags()],
-        "extra_link_args": kwargs.pop("extra_link_args", []),
-        **kwargs,
-    }
-
-
-register("attn_decode")
-register("attn_prefill")
-register("attn_paged_decode")
-register("rotary_emb")

csrc/kernels/attn_common.h

@@ -1,71 +1,77 @@
 #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_layout_policies.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 {
+    // Shape
     int batch;
     int q_head;
     int kv_head;
-    int q_len;
-    int kv_len;
     int head_dim;
-    int use_mask;
-    int causal_offset;   // -1 = non-causal; >=0 = absolute position of first Q token
-    int num_splits;
+    int q_len;   // Per-request in contiguous mode; total_q in paged mode.
+    int kv_len;  // Contiguous mode; paged mode uses kv_indptr.
+
+    // Attention behavior
     float scale;
+    // -1 = non-causal; >=0 = absolute position of first Q token
+    int causal_offset;
+    int use_mask;
 
-    // Q strides (element offsets for each dim — layout-agnostic)
-    int q_stride_b, q_stride_h, q_stride_l, q_stride_d;
-    // KV strides (K and V share the same layout — only base pointers differ)
-    int kv_stride_b, kv_stride_h, kv_stride_l, kv_stride_d;
-
-    // Mask: 2D [batch, kv_len], 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 T* __restrict__ q;
-    const T* __restrict__ k;
-    const T* __restrict__ v;
+    // pointers
+    const T* __restrict__ q_ptr;
+    const T* __restrict__ k_ptr;
+    const T* __restrict__ v_ptr;
+    const T* __restrict__ new_k_ptr;
+    const T* __restrict__ new_v_ptr;
+    T* __restrict__ o_ptr;
     const bool* __restrict__ mask;
 
-    T* __restrict__ o;
-    AT* __restrict__ o_part;
-    AT* __restrict__ ml_part;
-};
+    // strides
+    int q_b_stride;
+    int q_h_stride;
+    int q_l_stride;
+    int q_d_stride;
 
-template<typename T, typename AT = float>
-struct PagedAttentionParams {
-    int batch;
-    int q_head;
-    int kv_head;
-    int q_len;
-    int kv_len;
-    int head_dim;
-    int use_mask;
-    int causal_offset;
-    float scale;
+    int kv_b_stride;
+    int kv_h_stride;
+    int kv_l_stride;
+    int kv_d_stride;
 
-    int num_splits;
-    int page_size;
-    int max_pages;
+    int new_kv_b_stride;
+    int new_kv_h_stride;
 
-    // Q strides (layout-agnostic)
-    int q_stride_b, q_stride_h, q_stride_l, q_stride_d;
-
-    // Mask strides (2D, 3D, or 4D)
     int mask_b_stride;
     int mask_h_stride;
-    int mask_q_stride;
+    int mask_l_stride;
 
-    const T* __restrict__ q;
-    const T* __restrict__ k_cache;
-    const T* __restrict__ v_cache;
-    const bool* __restrict__ mask;
-    const int64_t* __restrict__ page_table;
+    // Paged K/V addressing
+    const int* __restrict__ req_to_token;         // [num_reqs, max_context_len]
+    const int* __restrict__ req_pool_indices;     // [batch]
+    const int* __restrict__ kv_indptr;             // [batch + 1]
+    const int* __restrict__ qo_indptr;             // [batch + 1] or nullptr for decode
+    const int* __restrict__ q_tile_to_batch;        // [num_q_tiles], prefill only
+    const int* __restrict__ q_tile_to_index;        // [num_q_tiles], prefill only
+    int num_q_tiles;
+    int max_context_len; // req_to_token stride (dim 1)
 
-    T* __restrict__ o;
+    // Decode split-KV workspace
+    int num_splits;
     AT* __restrict__ o_part;
     AT* __restrict__ ml_part;
+
 };

csrc/kernels/attn_decode.cu

@@ -8,19 +8,43 @@ torch::Tensor attn_decode(
     c10::optional<torch::Tensor> mask,
     int64_t causal_offset,
     double scale,
-    int64_t layout
+    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 == 1) ? O.transpose(1, 2) : O;
-    p.o = (bf16*)O_view.data_ptr();
+    auto O_view = (layout == BLHD) ? O.transpose(1, 2) : O;
+    p.o_ptr = (bf16*)O_view.data_ptr();
 
-    alloc_split_partials(p);
-    DISPATCH_HEAD_DIM(p.head_dim, dispatch_decode, p);
+    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;
+        int64_t ml_needed = (int64_t)p.batch * p.q_head * MAX_SPLITS * 2;
+        TORCH_CHECK(o_part_buf->numel() >= o_needed,
+                     "o_part_buf too small: need ", o_needed, " got ", o_part_buf->numel());
+        TORCH_CHECK(ml_part_buf->numel() >= ml_needed,
+                     "ml_part_buf too small: need ", ml_needed, " got ", ml_part_buf->numel());
+        TORCH_CHECK(o_part_buf->is_cuda() && ml_part_buf->is_cuda(),
+                     "split buffers must be CUDA tensors");
+        TORCH_CHECK(o_part_buf->is_contiguous() && ml_part_buf->is_contiguous(),
+                     "split buffers must be contiguous");
+        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;
 }
 
@@ -32,6 +56,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
         py::arg("mask") = py::none(),
         py::arg("causal_offset") = -1,
         py::arg("scale") = 0.0,
-        py::arg("layout") = 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)");
 }

csrc/kernels/attn_decode_split_kv.cuh

@@ -2,10 +2,16 @@
 #include <cuda_bf16.h>
 #include <float.h>
 #include "attn_common.h"
+#include "attn_layout_policies.cuh"
 #include "attn_warp_utils.cuh"
 constexpr int DC_CHUNK = 64;
 
-template <int HEAD_DIM, bool IsCausal, bool HasMask>
+// 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;
@@ -15,40 +21,46 @@ __global__ void attn_decode_split_kv_kernel(AttentionParams<bf16> p) {
     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 = batch * p.q_stride_b + q_head * p.q_stride_h
-              + lane * hd_per_thread * p.q_stride_d;
+    int q_off = KV::q_decode_base(p, batch, q_head)
+              + lane * hd_per_thread * p.q_d_stride;
     for (int i = 0; i < hd_per_thread; i++)
-        q_reg[i] = __bfloat162float(p.q[q_off + i * p.q_stride_d]);
+        q_reg[i] = __bfloat162float(p.q_ptr[q_off + i * p.q_d_stride]);
 
-    // KV: [batch, kv_head, kv_len, head_dim] — stride-based base
-    int kv_base = batch * p.kv_stride_b + kv_head * p.kv_stride_h;
     int mask_base = batch * p.mask_b_stride + q_head * p.mask_h_stride;
 
     float m = -FLT_MAX, d = 0.0f, acc_reg[8] = {0.0f};
 
-    extern __shared__ __align__(16) bf16 k_smem[];
+    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 = (p.kv_len + DC_CHUNK - 1) / DC_CHUNK;
+    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, p.kv_len - chunk_start);
+        int this_chunk = min(DC_CHUNK, seq_len - chunk_start);
 
-        // Load K into shared memory (gather from strided global)
+        // 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 kv_idx = chunk_start + s;
-            int g_off = kv_base + kv_idx * p.kv_stride_l + d_dim * p.kv_stride_d;
-            k_smem[i] = p.k[g_off];
+            int kc = chunk_start + s;
+            KVAddr a = KV::template decode_addr<1>(
+                p, kctx, batch, kv_head, kc, d_dim, true, 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();
 
@@ -65,7 +77,7 @@ __global__ void attn_decode_split_kv_kernel(AttentionParams<bf16> p) {
                     partial = -FLT_MAX;
             }
             if constexpr (IsCausal) {
-                if (kv_idx > p.causal_offset)
+                if (kv_idx >= KV::decode_attend_len(p, batch))
                     partial = -FLT_MAX;
             }
 
@@ -74,11 +86,10 @@ __global__ void attn_decode_split_kv_kernel(AttentionParams<bf16> p) {
             float beta  = __expf(partial - new_m);
             d = d * alpha + beta;
 
-            int v_off = kv_base + kv_idx * p.kv_stride_l
-                        + lane * hd_per_thread * p.kv_stride_d;
-            for (int i = 0; i < hd_per_thread; i++)
-                acc_reg[i] = fmaf(acc_reg[i], alpha,
-                                  __bfloat162float(p.v[v_off + i * p.kv_stride_d]) * 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();
@@ -98,6 +109,10 @@ __global__ void attn_decode_split_kv_kernel(AttentionParams<bf16> p) {
     }
 }
 
+// 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;
@@ -124,6 +139,6 @@ __global__ void attn_decode_combine_kernel(AttentionParams<bf16> p) {
     }
 
     float inv = (l > 1e-20f) ? (1.0f / l) : 0.0f;
-    int o_off = batch * p.q_stride_b + q_head * p.q_stride_h + d * p.q_stride_d;
-    p.o[o_off] = __float2bfloat16(acc * inv);
+    int o_off = KV::q_decode_base(p, batch, q_head) + d * p.q_d_stride;
+    p.o_ptr[o_off] = __float2bfloat16(acc * inv);
 }

csrc/kernels/attn_decode_split_kv_mma.cuh

@@ -2,19 +2,22 @@
 #include <cfloat>
 #include <cuda_bf16.h>
 #include "attn_common.h"
+#include "attn_layout_policies.cuh"
 #include "attn_mma_utils.cuh"
 #include "attn_warp_utils.cuh"
 
-// Split-K (FlashDecoding) tensor-core decode via GQA head-packing.
-// 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.
+// 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=32, WARPS=1, STAGES=<2 or 1>>.
-template <typename Traits, bool IsCausal, bool HasMask>
+// 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;
@@ -31,18 +34,21 @@ __global__ void attn_decode_split_kv_mma_kernel(AttentionParams<bf16> p) {
     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.
-    // stride_row = p.q_stride_h for decode (q_len=1).
-    const int q_base = batch * p.q_stride_b + q_head0 * p.q_stride_h;
+    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,
+    load_q_mma_frags<Traits::KD>(p.q_ptr + q_base, p.q_h_stride, p.q_d_stride,
                                   qra, qrb, va, vb, tid4, Qa);
 
     float Oacc[Traits::DN8][4];
@@ -51,13 +57,12 @@ __global__ void attn_decode_split_kv_mma_kernel(AttentionParams<bf16> p) {
         Oacc[j][0] = Oacc[j][1] = Oacc[j][2] = Oacc[j][3] = 0.0f;
     float m0 = -FLT_MAX, m1 = -FLT_MAX, l0 = 0.0f, l1 = 0.0f;
 
-    const int kv_base = batch * p.kv_stride_b + kv_head * p.kv_stride_h;
-    const int tiles_total = (p.kv_len + Traits::BC - 1) / Traits::BC;
+    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 ----
+    // ---- 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;
@@ -67,11 +72,14 @@ __global__ void attn_decode_split_kv_mma_kernel(AttentionParams<bf16> p) {
              i += Traits::NUM_THREADS * Traits::VEC) {
             int r = i / Traits::HEAD_DIM, d = i % Traits::HEAD_DIM;
             int kc = kv0 + r;
-            bool valid = kc < p.kv_len;
+            bool valid = kc < seq_len;
+            // All GQA passes consume new K/V directly. Only the first pass
+            // persists it, so no cross-block synchronization is required.
+            KVAddr a = KV::template decode_addr<Traits::VEC>(
+                p, kctx, batch, kv_head, kc, d, valid, pass == 0);
             int off = r * Traits::LD + swiz_col(d, r, Traits::SWIZ_MASK);
-            int g_off = kv_base + kc * p.kv_stride_l + d * p.kv_stride_d;
-            cp_async_16_pred(&dK[off], &p.k[g_off], valid);
-            cp_async_16_pred(&dV[off], &p.v[g_off], valid);
+            cp_async_16_pred(&dK[off], a.k, a.valid);
+            cp_async_16_pred(&dV[off], a.v, a.valid);
         }
         cp_async_commit();
     };
@@ -96,14 +104,17 @@ __global__ void attn_decode_split_kv_mma_kernel(AttentionParams<bf16> p) {
             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
-        int maxc = IsCausal ? min(p.kv_len, p.causal_offset + 1) : p.kv_len;
+        // 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);
+                                            p.mask_b_stride, p.mask_h_stride, p.mask_l_stride,
+                                            batch, q_head0 + gid, q_head0 + gid + 8,
+                                            p.mask,
+                                            va, vb,
+                                            Sacc, Oacc, m0, m1, l0, l1, lane);
 
         mma_pv_accumulate<Traits>(Sacc, bV, lane, Oacc);
     };
@@ -114,7 +125,10 @@ __global__ void attn_decode_split_kv_mma_kernel(AttentionParams<bf16> p) {
             load_tile(ti_begin + i, i);
 
         for (int it = 0; it < ntiles; it++) {
-            cp_async_wait_group<STAGES - 1>();
+            if (it + 1 == ntiles)
+                cp_async_wait_group<0>();
+            else
+                cp_async_wait_group<STAGES - 1>();
             __syncwarp();
             process_tile(it, it & (STAGES - 1));
             __syncwarp();

csrc/kernels/attn_dispatchers.cuh

@@ -1,81 +1,153 @@
 #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_layout_policies.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_layout_policies.cuh"
 #include "attn_prefill_split_q.cuh"
 #include "attn_decode_split_kv.cuh"
-#include "attn_paged_decode_split_kv.cuh"
 #ifndef ASTRAI_NO_MMA
 #include "attn_prefill_split_q_mma.cuh"
 #include "attn_decode_split_kv_mma.cuh"
-#include "attn_paged_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 sm_count = 0;
-    cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, 0);
-    int n = (2 * sm_count + base_blocks - 1) / base_blocks;
+                               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
+// Prefill launchers (KV selects ContigKV or PagedKV addressing)
 // ======================================================================
 
 #ifndef ASTRAI_NO_MMA
-template <int HEAD_DIM, bool IsCausal, bool HasMask>
-static inline void launch_prefill_mma(AttentionParams<bf16>& p) {
-    constexpr int WARPS = 4;
-    constexpr int BC = (HEAD_DIM <= 128) ? 32 : 16;
-    using Traits = KernelTraits<HEAD_DIM, BC, WARPS, 2>;
-    dim3 grid((p.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, IsCausal, HasMask><<<grid, block>>>(p);
-}
+template <int BC_>
+struct PrefillKernelConfig {
+    static constexpr int BC = BC_;
+    static constexpr int WARPS = 4;
+    static constexpr int STAGES = 2;
+};
+
+// Compile-time configuration map shared by contiguous and paged prefill.
+// Unsupported head dimensions intentionally have no mapping.
+template <int HEAD_DIM, bool IsCausal>
+struct PrefillConfigMap;
+
+template <> struct PrefillConfigMap<32, false>  : PrefillKernelConfig<32> {};
+template <> struct PrefillConfigMap<32, true>   : PrefillKernelConfig<64> {};
+template <> struct PrefillConfigMap<64, false>  : PrefillKernelConfig<32> {};
+template <> struct PrefillConfigMap<64, true>   : PrefillKernelConfig<64> {};
+template <> struct PrefillConfigMap<128, false> : PrefillKernelConfig<32> {};
+template <> struct PrefillConfigMap<128, true>  : PrefillKernelConfig<32> {};
+template <> struct PrefillConfigMap<256, false> : PrefillKernelConfig<16> {};
+template <> struct PrefillConfigMap<256, true>  : PrefillKernelConfig<16> {};
+
+template <typename QSchedule, typename KV>
+struct PrefillLauncherMMA {
+    template <int HEAD_DIM, bool IsCausal, bool HasMask>
+    static void launch(AttentionParams<bf16>& p, cudaStream_t stream) {
+        using Config = PrefillConfigMap<HEAD_DIM, IsCausal>;
+        using Traits = KernelTraits<HEAD_DIM, Config::BC, Config::WARPS, Config::STAGES>;
+        constexpr int ROWS = Traits::BR * Config::WARPS;
+        dim3 grid(QSchedule::host_q_blocks(p, ROWS), p.q_head,
+                  QSchedule::host_grid_batch(p));
+        dim3 block(Traits::NUM_THREADS);
+        attn_prefill_split_q_mma_kernel<Traits, QSchedule, KV, IsCausal, HasMask>
+            <<<grid, block, 0, stream>>>(p);
+    }
+};
 #endif
 
-template <int HEAD_DIM, bool IsCausal, bool HasMask>
-static inline void launch_prefill_scalar(AttentionParams<bf16>& p) {
-    constexpr int G = 8, ROWS = 32, P_BC = 32;
-    dim3 grid((p.q_len + ROWS - 1) / ROWS, p.q_head, p.batch);
-    dim3 block(G, ROWS);
-    attn_prefill_split_q_kernel_t<HEAD_DIM, G, ROWS, P_BC, IsCausal, HasMask><<<grid, block>>>(p);
-}
+template <typename QSchedule, typename KV>
+struct PrefillLauncherScalar {
+    template <int HEAD_DIM, bool IsCausal, bool HasMask>
+    static void launch(AttentionParams<bf16>& p, cudaStream_t stream) {
+        constexpr int G = (HEAD_DIM == 32) ? 4 : 8, ROWS = 64, P_BC = 32;
+        dim3 grid(QSchedule::host_q_blocks(p, ROWS), p.q_head,
+                  QSchedule::host_grid_batch(p));
+        dim3 block(G, ROWS);
+        attn_prefill_split_q_kernel_t<HEAD_DIM, QSchedule, KV, G, ROWS, P_BC,
+                                      IsCausal, HasMask>
+            <<<grid, block, 0, stream>>>(p);
+    }
+};
 
 template <int HEAD_DIM>
-static inline void dispatch_prefill(AttentionParams<bf16>& p) {
+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
-    if (is_causal) {
-        if (has_mask)      launch_prefill_mma<HEAD_DIM, true, true>(p);
-        else               launch_prefill_mma<HEAD_DIM, true, false>(p);
-    } else {
-        if (has_mask)      launch_prefill_mma<HEAD_DIM, false, true>(p);
-        else               launch_prefill_mma<HEAD_DIM, false, false>(p);
-    }
+    using Launcher = PrefillLauncherMMA<DenseQSchedule, ContigKV>;
+    DISPATCH_CAUSAL_MASK(is_causal, has_mask,
+                         Launcher::template launch,
+                         HEAD_DIM, p, stream);
 #else
-    if (is_causal) {
-        if (has_mask)      launch_prefill_scalar<HEAD_DIM, true, true>(p);
-        else               launch_prefill_scalar<HEAD_DIM, true, false>(p);
-    } else {
-        if (has_mask)      launch_prefill_scalar<HEAD_DIM, false, true>(p);
-        else               launch_prefill_scalar<HEAD_DIM, false, false>(p);
-    }
+    using Launcher = PrefillLauncherScalar<DenseQSchedule, ContigKV>;
+    DISPATCH_CAUSAL_MASK(is_causal, has_mask,
+                         Launcher::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
+    using Launcher = PrefillLauncherMMA<PackedQSchedule, PagedKV>;
+    DISPATCH_CAUSAL_MASK(is_causal, has_mask,
+                         Launcher::template launch,
+                         HEAD_DIM, p, stream);
+#else
+    using Launcher = PrefillLauncherScalar<PackedQSchedule, PagedKV>;
+    DISPATCH_CAUSAL_MASK(is_causal, has_mask,
+                         Launcher::template launch,
+                         HEAD_DIM, p, stream);
 #endif
 }
 
 // ======================================================================
-// Decode
+// Decode launchers (KV selects ContigKV or PagedKV addressing)
 // ======================================================================
 
 #ifndef ASTRAI_NO_MMA
@@ -83,113 +155,79 @@ static inline void dispatch_prefill(AttentionParams<bf16>& p) {
 // 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 <int HEAD_DIM, bool IsCausal, bool HasMask>
-static inline void launch_decode_mma(AttentionParams<bf16>& p, int group_size) {
-    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 tiles_total = (p.kv_len + BC - 1) / BC;
-    p.num_splits = compute_num_splits(p.batch * p.kv_head, 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, IsCausal, HasMask><<<grid, 32>>>(p);
-}
+template <typename KV>
+struct DecodeLauncherMMA {
+    template <int HEAD_DIM, bool IsCausal, bool HasMask>
+    static void launch(AttentionParams<bf16>& p, 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 <int HEAD_DIM, bool IsCausal, bool HasMask>
-static inline void launch_decode_scalar(AttentionParams<bf16>& p, int group_size) {
-    int chunks_total = (p.kv_len + DC_CHUNK - 1) / DC_CHUNK;
-    p.num_splits = compute_num_splits(p.batch * p.kv_head, chunks_total);
-    size_t smem = 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, IsCausal, HasMask><<<grid, block, smem>>>(p);
+template <typename KV>
+struct DecodeLauncherScalar {
+    template <int HEAD_DIM, bool IsCausal, bool HasMask>
+    static void launch(AttentionParams<bf16>& p, 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 group_size = p.q_head / p.kv_head;
+        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);
+        cudaFuncSetAttribute(
+            attn_decode_split_kv_kernel<HEAD_DIM, KV, IsCausal, HasMask>,
+            cudaFuncAttributeMaxDynamicSharedMemorySize,
+            smem);
+        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);
+
+#ifndef ASTRAI_NO_MMA
+    DISPATCH_CAUSAL_MASK(is_causal, has_mask,
+                         DecodeLauncherMMA<ContigKV>::template launch,
+                         HEAD_DIM, p, stream);
+#else
+    DISPATCH_CAUSAL_MASK(is_causal, has_mask,
+                         DecodeLauncherScalar<ContigKV>::template launch,
+                         HEAD_DIM, p, 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_decode(AttentionParams<bf16>& p) {
+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
-    if (is_causal) {
-        if (has_mask)      launch_decode_mma<HEAD_DIM, true, true>(p, group_size);
-        else               launch_decode_mma<HEAD_DIM, true, false>(p, group_size);
-    } else {
-        if (has_mask)      launch_decode_mma<HEAD_DIM, false, true>(p, group_size);
-        else               launch_decode_mma<HEAD_DIM, false, false>(p, group_size);
-    }
+    DISPATCH_CAUSAL_MASK(is_causal, has_mask,
+                         DecodeLauncherMMA<PagedKV>::template launch,
+                         HEAD_DIM, p, stream);
 #else
-    if (is_causal) {
-        if (has_mask)      launch_decode_scalar<HEAD_DIM, true, true>(p, group_size);
-        else               launch_decode_scalar<HEAD_DIM, true, false>(p, group_size);
-    } else {
-        if (has_mask)      launch_decode_scalar<HEAD_DIM, false, true>(p, group_size);
-        else               launch_decode_scalar<HEAD_DIM, false, false>(p, group_size);
-    }
+    DISPATCH_CAUSAL_MASK(is_causal, has_mask,
+                         DecodeLauncherScalar<PagedKV>::template launch,
+                         HEAD_DIM, p, stream);
 #endif
 
-    attn_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim>>>(p);
-}
-
-// ======================================================================
-// Paged Decode
-// ======================================================================
-
-#ifndef ASTRAI_NO_MMA
-template <int HEAD_DIM, bool IsCausal, bool HasMask>
-static inline void launch_paged_decode_mma(PagedAttentionParams<bf16>& p, int group_size) {
-    int G = p.q_head / p.kv_head;
-    constexpr int MAX_G = 16;
-    constexpr int BC = 16;
-    int num_passes = (G + MAX_G - 1) / MAX_G;
-    int tiles_total = (p.kv_len + BC - 1) / BC;
-    p.num_splits = compute_num_splits(p.batch * p.kv_head, 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);
-    paged_attn_decode_split_kv_mma_kernel<Traits, IsCausal, HasMask> <<<grid, 32>>>(p);
-}
-#endif
-
-template <int HEAD_DIM, bool IsCausal, bool HasMask>
-static inline void launch_paged_decode_scalar(PagedAttentionParams<bf16>& p, int group_size) {
-    int chunks_total = (p.kv_len + PDC_CHUNK - 1) / PDC_CHUNK;
-    p.num_splits = compute_num_splits(p.batch * p.kv_head, chunks_total);
-    size_t smem = PDC_CHUNK * p.head_dim * sizeof(bf16);
-    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);
-    paged_attn_decode_split_kv_kernel<HEAD_DIM, IsCausal, HasMask><<<grid, block, smem>>>(p);
-}
-
-template <int HEAD_DIM>
-static inline void dispatch_paged_decode(PagedAttentionParams<bf16>& p) {
-    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
-    if (is_causal) {
-        if (has_mask)      launch_paged_decode_mma<HEAD_DIM, true, true>(p, group_size);
-        else               launch_paged_decode_mma<HEAD_DIM, true, false>(p, group_size);
-    } else {
-        if (has_mask)      launch_paged_decode_mma<HEAD_DIM, false, true>(p, group_size);
-        else               launch_paged_decode_mma<HEAD_DIM, false, false>(p, group_size);
-    }
-#else
-    if (is_causal) {
-        if (has_mask)      launch_paged_decode_scalar<HEAD_DIM, true, true>(p, group_size);
-        else               launch_paged_decode_scalar<HEAD_DIM, true, false>(p, group_size);
-    } else {
-        if (has_mask)      launch_paged_decode_scalar<HEAD_DIM, false, true>(p, group_size);
-        else               launch_paged_decode_scalar<HEAD_DIM, false, false>(p, group_size);
-    }
-#endif
-
-    paged_attn_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim>>>(p);
+    attn_decode_combine_kernel<PagedKV><<<p.batch * p.q_head, p.head_dim, 0, stream>>>(p);
 }

csrc/kernels/attn_entry_utils.cuh

@@ -8,14 +8,14 @@
 using bf16 = __nv_bfloat16;
 
 // Dispatch head_dim: shared macro — avoids C++20 lambda template syntax.
-// Usage: DISPATCH_HEAD_DIM(hd, fn, arg)
-//   Expands to: fn<32>(arg); fn<64>(arg); etc.
-#define DISPATCH_HEAD_DIM(hd, fn, arg) \
+// 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>(arg); break; \
-        case 64:  fn<64>(arg); break; \
-        case 128: fn<128>(arg); break; \
-        case 256: fn<256>(arg); break; \
+        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)"); \
@@ -37,15 +37,15 @@ inline void alloc_split_partials(P& p) {
 // ---- Shared Q-dims + strides extraction ----
 template <typename P>
 inline void extract_q_dims_and_strides(torch::Tensor& q, int64_t layout, P& p) {
-    if (layout == 1) q = q.transpose(1, 2);
+    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);
+    p.q_b_stride = (int)q.stride(0);
+    p.q_h_stride = (int)q.stride(1);
+    p.q_l_stride = (int)q.stride(2);
+    p.q_d_stride = (int)q.stride(3);
 }
 
 // ---- Shared mask packing ----
@@ -63,17 +63,17 @@ inline void pack_mask(const c10::optional<torch::Tensor>& mask, P& p) {
         if (m.dim() == 2) {
             p.mask_b_stride = (int)m.stride(0);
             p.mask_h_stride = 0;
-            p.mask_q_stride = 0;
+            p.mask_l_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);
+            p.mask_l_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);
+            p.mask_l_stride = (m.size(2) == 1) ? 0 : (int)m.stride(2);
         } else {
             TORCH_CHECK(false, "mask must be 2D, 3D, or 4D");
         }
@@ -82,7 +82,7 @@ inline void pack_mask(const c10::optional<torch::Tensor>& mask, P& p) {
         p.mask = nullptr;
         p.mask_b_stride = 0;
         p.mask_h_stride = 0;
-        p.mask_q_stride = 0;
+        p.mask_l_stride = 0;
     }
 }
 
@@ -106,82 +106,253 @@ inline void attn_pack_params(
     TORCH_CHECK(v.dtype() == torch::kBFloat16);
     TORCH_CHECK(k.sizes() == v.sizes(), "K and V must have identical shapes");
     TORCH_CHECK(q.dim() == 4 && k.dim() == 4, "Q/K/V must be 4D");
-
     extract_q_dims_and_strides(q, layout, p);
 
-    if (layout == 1) k = k.transpose(1, 2), v = v.transpose(1, 2);
+    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(p.q_head % p.kv_head == 0,
+                "q_head must be divisible by kv_head");
     TORCH_CHECK(k.size(3) == p.head_dim, "K/V head_dim must match Q");
+    TORCH_CHECK(q.stride(3) == 1 && k.stride(3) == 1 && v.stride(3) == 1,
+                "Q/K/V head_dim must be contiguous");
 
-    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.kv_b_stride = (int)k.stride(0);
+    p.kv_h_stride = (int)k.stride(1);
+    p.kv_l_stride = (int)k.stride(2);
+    p.kv_d_stride = (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.q_ptr = (const T*)q.data_ptr();
+    p.k_ptr = (const T*)k.data_ptr();
+    p.v_ptr = (const T*)v.data_ptr();
+    p.new_k_ptr = nullptr;
+    p.new_v_ptr = nullptr;
+    p.o_ptr = nullptr;
     p.o_part = nullptr;
     p.ml_part = nullptr;
 
     pack_mask(mask, p);
 }
 
-// ---- attn_pack_paged_params ----
+// ---- 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_params(
+inline void attn_pack_paged_decode_params(
     torch::Tensor q,
-    torch::Tensor page_table,
     torch::Tensor k_cache,
     torch::Tensor v_cache,
-    int64_t page_size,
-    int64_t kv_len,
+    torch::Tensor req_to_token,
+    torch::Tensor req_pool_indices,
+    torch::Tensor kv_indptr,
+    const c10::optional<torch::Tensor>& new_k,
+    const c10::optional<torch::Tensor>& new_v,
     c10::optional<torch::Tensor> mask,
     int64_t causal_offset,
     double scale,
-    int64_t layout,
-    PagedAttentionParams<T>& p
+    AttentionParams<T>& p
 ) {
     const at::cuda::OptionalCUDAGuard device_guard(device_of(q));
 
-    TORCH_CHECK(q.is_cuda() && page_table.is_cuda() && k_cache.is_cuda() && v_cache.is_cuda());
+    TORCH_CHECK(q.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(page_table.dtype() == torch::kLong, "page_table must be int64");
-    TORCH_CHECK(k_cache.sizes() == v_cache.sizes(), "k_cache and v_cache must have identical shapes");
+    TORCH_CHECK(req_to_token.dtype() == torch::kInt32, "req_to_token must be int32");
+    TORCH_CHECK(req_pool_indices.dtype() == torch::kInt32,
+                "req_pool_indices must be int32");
+    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]");
 
-    extract_q_dims_and_strides(q, layout, p);
-
-    p.kv_head = (int)k_cache.size(2);
-    p.kv_len = (int)kv_len;
-    p.page_size = (int)page_size;
-    p.max_pages = (int)page_table.size(1);
-
-    TORCH_CHECK(q.size(2) == 1, "Q seq_len must be 1 (decode)");
+    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(q.stride(2) == 1 && k_cache.stride(2) == 1 && v_cache.stride(2) == 1,
+                "Q/K/V head_dim must be contiguous");
     TORCH_CHECK(p.head_dim % 32 == 0, "head_dim must be multiple of 32");
-    TORCH_CHECK(k_cache.size(1) == page_size,
-                "k_cache dim 1 must equal page_size, got ",
-                k_cache.size(1), " vs ", page_size);
+    TORCH_CHECK(p.q_head % p.kv_head == 0, "q_head must be divisible by kv_head");
+
+    p.q_l_stride = (int)q.stride(0);
+    p.q_h_stride = (int)q.stride(1);
+    p.q_d_stride = (int)q.stride(2);
+
+    p.k_ptr = (const T*)k_cache.data_ptr();
+    p.v_ptr = (const T*)v_cache.data_ptr();
+    p.q_ptr = (const T*)q.data_ptr();
+    p.req_to_token = req_to_token.data_ptr<int>();
+    p.req_pool_indices = req_pool_indices.data_ptr<int>();
+    p.kv_indptr = kv_indptr.data_ptr<int>();
+    p.qo_indptr = nullptr;
+    p.max_context_len = (int)req_to_token.size(1);
+
+    TORCH_CHECK(new_k.has_value() == new_v.has_value(),
+                "new_k and new_v must be provided together");
+    if (new_k.has_value()) {
+        auto nk = new_k.value();
+        auto nv = new_v.value();
+        TORCH_CHECK(nk.is_cuda() && nv.is_cuda(), "new K/V must be CUDA tensors");
+        TORCH_CHECK(nk.dtype() == torch::kBFloat16 && nv.dtype() == torch::kBFloat16,
+                    "new K/V must be bf16");
+        TORCH_CHECK(nk.dim() == 3 && nv.dim() == 3,
+                    "new K/V must be 3D [batch, kv_head, head_dim]");
+        TORCH_CHECK(nk.sizes() == nv.sizes(), "new K and V must have identical shapes");
+        TORCH_CHECK(nk.strides() == nv.strides(),
+                    "new K and V must have identical strides");
+        TORCH_CHECK(nk.size(0) == p.batch && nk.size(1) == p.kv_head
+                    && nk.size(2) == p.head_dim, "new K/V shape mismatch");
+        TORCH_CHECK(nk.stride(2) == 1 && nv.stride(2) == 1,
+                    "new K/V head_dim must be contiguous");
+        p.new_k_ptr = (const T*)nk.data_ptr();
+        p.new_v_ptr = (const T*)nv.data_ptr();
+        p.new_kv_b_stride = (int)nk.stride(0);
+        p.new_kv_h_stride = (int)nk.stride(1);
+    } else {
+        p.new_k_ptr = nullptr;
+        p.new_v_ptr = nullptr;
+        p.new_kv_b_stride = p.new_kv_h_stride = 0;
+    }
 
     p.causal_offset = (int)causal_offset;
     p.use_mask = (mask.has_value() && mask.value().defined()) ? 1 : 0;
     p.scale = (scale > 0.0) ? (float)scale : 1.0f / sqrtf((float)p.head_dim);
 
-    p.page_table = page_table.data_ptr<int64_t>();
-    p.k_cache = (const T*)k_cache.data_ptr();
-    p.v_cache = (const T*)v_cache.data_ptr();
-    p.q = (const T*)q.data_ptr();
-    p.o = nullptr;
+    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_l_stride = 0;
+        p.mask = m.data_ptr<bool>();
+    } else {
+        p.mask = nullptr;
+        p.mask_b_stride = 0;
+        p.mask_h_stride = 0;
+        p.mask_l_stride = 0;
+    }
+
+    p.o_ptr = 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,
+    torch::Tensor q_tile_to_batch,
+    torch::Tensor q_tile_to_index,
+    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());
+    TORCH_CHECK(kv_indptr.is_cuda() && qo_indptr.is_cuda());
+    TORCH_CHECK(q_tile_to_batch.is_cuda() && q_tile_to_index.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::kInt32, "req_to_token must be int32");
+    TORCH_CHECK(req_pool_indices.dtype() == torch::kInt32,
+                "req_pool_indices must be int32");
+    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(q_tile_to_batch.dtype() == torch::kInt32,
+                "q_tile_to_batch must be int32");
+    TORCH_CHECK(q_tile_to_index.dtype() == torch::kInt32,
+                "q_tile_to_index 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.q_len = (int)q.size(0);
+    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(q.stride(2) == 1 && k_cache.stride(2) == 1 && v_cache.stride(2) == 1,
+                "Q/K/V head_dim must be contiguous");
+    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]");
+    TORCH_CHECK(q_tile_to_batch.dim() == 1 && q_tile_to_index.dim() == 1,
+                "Q tile mappings must be 1D");
+    TORCH_CHECK(q_tile_to_batch.size(0) == q_tile_to_index.size(0),
+                "Q tile mappings must have equal length");
+
+    p.q_l_stride = (int)q.stride(0);
+    p.q_h_stride = (int)q.stride(1);
+    p.q_d_stride = (int)q.stride(2);
+
+    p.k_ptr = (const T*)k_cache.data_ptr();
+    p.v_ptr = (const T*)v_cache.data_ptr();
+    p.new_k_ptr = nullptr;
+    p.new_v_ptr = nullptr;
+    p.q_ptr = (const T*)q.data_ptr();
+    p.req_to_token = req_to_token.data_ptr<int>();
+    p.req_pool_indices = req_pool_indices.data_ptr<int>();
+    p.kv_indptr = kv_indptr.data_ptr<int>();
+    p.qo_indptr = qo_indptr.data_ptr<int>();
+    p.q_tile_to_batch = q_tile_to_batch.data_ptr<int>();
+    p.q_tile_to_index = q_tile_to_index.data_ptr<int>();
+    p.num_q_tiles = (int)q_tile_to_batch.size(0);
+    p.max_context_len = (int)req_to_token.size(1);
+
+    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_l_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) > 0 && m.size(2) <= p.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_l_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_l_stride = 0;
+    }
+    p.scale = (scale > 0.0) ? (float)scale : 1.0f / sqrtf((float)p.head_dim);
+
+    p.o_ptr = nullptr;
     p.o_part = nullptr;
     p.ml_part = nullptr;
-
-    pack_mask(mask, p);
 }

csrc/kernels/attn_layout_policies.cuh

@@ -0,0 +1,255 @@
+#pragma once
+#include <cuda_bf16.h>
+#include "attn_common.h"
+
+// ============================================================================
+// Attention layout policies keep Q scheduling independent from K/V storage.
+// DenseQSchedule / PackedQSchedule map blocks to Q tiles; ContigKV / PagedKV
+// resolve logical K/V positions to physical addresses. This lets the shared
+// kernels compose Q layout and K/V storage without coupling the two concerns.
+//
+//   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_b_stride, 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_l_stride.
+//
+// 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.
+// ============================================================================
+
+#define HOST_FORCEINLINE static __host__ __forceinline__
+#define DEVICE_FORCEINLINE static __device__ __forceinline__
+#define HOST_DEV_FORCEINLINE static __host__ __device__ __forceinline__
+
+using bf16 = __nv_bfloat16;
+
+// ============================================================================
+// Q scheduling policies
+//
+// Map CUDA blocks to request-local Q tiles independently of K/V storage.
+// Dense tensors encode the request in blockIdx.z; packed ragged tensors use
+// a compact precomputed work map indexed by blockIdx.x.
+// ============================================================================
+
+struct DenseQSchedule {
+    HOST_FORCEINLINE int host_q_blocks(
+        const AttentionParams<bf16>& p, int rows) {
+        return (p.q_len + rows - 1) / rows;
+    }
+
+    HOST_FORCEINLINE int host_grid_batch(
+        const AttentionParams<bf16>& p) {
+        return p.batch;
+    }
+
+    DEVICE_FORCEINLINE void map_block(
+        const AttentionParams<bf16>&, int& batch, int& q_tile) {
+        batch = blockIdx.z;
+        q_tile = blockIdx.x;
+    }
+
+    DEVICE_FORCEINLINE int q_len(
+        const AttentionParams<bf16>& p, int) {
+        return p.q_len;
+    }
+
+    DEVICE_FORCEINLINE int q_base(
+        const AttentionParams<bf16>& p, int batch, int q_head) {
+        return batch * p.q_b_stride + q_head * p.q_h_stride;
+    }
+};
+
+struct PackedQSchedule {
+    HOST_FORCEINLINE int host_q_blocks(
+        const AttentionParams<bf16>& p, int) {
+        return p.num_q_tiles;
+    }
+
+    HOST_FORCEINLINE int host_grid_batch(
+        const AttentionParams<bf16>&) {
+        return 1;
+    }
+
+    DEVICE_FORCEINLINE void map_block(
+        const AttentionParams<bf16>& p, int& batch, int& q_tile) {
+        batch = p.q_tile_to_batch[blockIdx.x];
+        q_tile = p.q_tile_to_index[blockIdx.x];
+    }
+
+    DEVICE_FORCEINLINE int q_len(
+        const AttentionParams<bf16>& p, int batch) {
+        return p.qo_indptr[batch + 1] - p.qo_indptr[batch];
+    }
+
+    DEVICE_FORCEINLINE int q_base(
+        const AttentionParams<bf16>& p, int batch, int q_head) {
+        return p.qo_indptr[batch] * p.q_l_stride + q_head * p.q_h_stride;
+    }
+};
+
+// Hoisted per-(batch, kv_head) addressing context.
+struct KVContext {
+    int kv_base;          // contig: batch*kv_b_stride + kv_head*kv_h_stride
+    int 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_FORCEINLINE int host_kv_len(const AttentionParams<bf16>& p) {
+        return p.kv_len;
+    }
+
+    // decode: same offset (q_len == 1, so there is no row stride component)
+    DEVICE_FORCEINLINE int q_decode_base(
+        const AttentionParams<bf16>& p, int batch, int q_head) {
+        return batch * p.q_b_stride + q_head * p.q_h_stride;
+    }
+
+    DEVICE_FORCEINLINE int kv_len(const AttentionParams<bf16>& p, int) {
+        return p.kv_len;
+    }
+    DEVICE_FORCEINLINE int causal_offset(
+        const AttentionParams<bf16>& p, int, int) {
+        return p.causal_offset;
+    }
+    // decode: exclusive bound of the single query's attend range
+    DEVICE_FORCEINLINE int decode_attend_len(const AttentionParams<bf16>& p, int) {
+        return (p.kv_len < p.causal_offset + 1) ? p.kv_len : (p.causal_offset + 1);
+    }
+
+    template <int HEAD_DIM>
+    DEVICE_FORCEINLINE KVContext make_ctx(
+        const AttentionParams<bf16>& p, int batch, int kv_head) {
+        KVContext c = {};
+        c.kv_base = batch * p.kv_b_stride + kv_head * p.kv_h_stride;
+        return c;
+    }
+    DEVICE_FORCEINLINE int resolve_token(
+        const AttentionParams<bf16>& p, const KVContext& c, int kc, bool valid) {
+        return valid ? kc : -1;
+    }
+    DEVICE_FORCEINLINE KVAddr kv_addr_from_token(
+        const AttentionParams<bf16>& p, const KVContext& c, int token, int d) {
+        const bool valid = token >= 0;
+        const int safe_token = valid ? token : 0;
+        const int64_t gmem_off = (int64_t)c.kv_base
+            + (int64_t)safe_token * p.kv_l_stride
+            + (int64_t)d * p.kv_d_stride;
+        return {&p.k_ptr[gmem_off], &p.v_ptr[gmem_off], valid};
+    }
+
+    template <int VEC>
+    DEVICE_FORCEINLINE KVAddr decode_addr(
+        const AttentionParams<bf16>& p, const KVContext& c,
+        int, int, int kc, int d, bool valid, bool) {
+        int token = resolve_token(p, c, kc, valid);
+        return kv_addr_from_token(p, c, token, d);
+    }
+};
+
+// ---- Paged (SGLang-style flat pool) K/V ----
+struct PagedKV {
+    static constexpr bool kPaged = true;
+
+    HOST_FORCEINLINE int host_kv_len(const AttentionParams<bf16>& p) {
+        return p.max_context_len;
+    }
+
+    // decode: Q is [batch, q_head, head_dim], so batch is the outer row
+    DEVICE_FORCEINLINE int q_decode_base(
+        const AttentionParams<bf16>& p, int batch, int q_head) {
+        return batch * p.q_l_stride + q_head * p.q_h_stride;
+    }
+
+    DEVICE_FORCEINLINE int kv_len(const AttentionParams<bf16>& p, int batch) {
+        return p.kv_indptr[batch + 1] - p.kv_indptr[batch];
+    }
+    DEVICE_FORCEINLINE int causal_offset(
+        const AttentionParams<bf16>& p, int batch, int q_len) {
+        return kv_len(p, batch) - q_len;
+    }
+    // decode: the query is the last token, so [0, seq_len) IS its causal range
+    DEVICE_FORCEINLINE int decode_attend_len(const AttentionParams<bf16>& p, int batch) {
+        return kv_len(p, batch);
+    }
+
+    template <int HEAD_DIM>
+    DEVICE_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;
+    }
+    DEVICE_FORCEINLINE int resolve_token(
+        const AttentionParams<bf16>& p, const KVContext& c, int kc, bool valid) {
+        return valid ? p.req_to_token[c.req_idx * c.rtt_stride + kc] : -1;
+    }
+    DEVICE_FORCEINLINE KVAddr kv_addr_from_token(
+        const AttentionParams<bf16>& p, const KVContext& c, int slot, int d) {
+        const bool valid = slot >= 0;
+        const int safe_slot = valid ? slot : 0;
+        const int64_t gmem_off = (int64_t)safe_slot * c.pool_stride + c.head_off + d;
+        return {&p.k_ptr[gmem_off], &p.v_ptr[gmem_off], valid};
+    }
+
+    DEVICE_FORCEINLINE KVAddr new_kv_addr(
+        const AttentionParams<bf16>& p, int batch, int kv_head, int d) {
+        const int64_t off = (int64_t)batch * p.new_kv_b_stride
+            + (int64_t)kv_head * p.new_kv_h_stride + d;
+        return {&p.new_k_ptr[off], &p.new_v_ptr[off], true};
+    }
+
+    DEVICE_FORCEINLINE void store_new_kv(
+        const AttentionParams<bf16>& p, const KVContext& c,
+        int seq_len, int d, const KVAddr& src) {
+        int slot = resolve_token(p, c, seq_len - 1, true);
+        const int64_t off = (int64_t)slot * c.pool_stride + c.head_off + d;
+        const_cast<bf16*>(p.k_ptr)[off] = *reinterpret_cast<const bf16*>(src.k);
+        const_cast<bf16*>(p.v_ptr)[off] = *reinterpret_cast<const bf16*>(src.v);
+    }
+
+    template <int VEC>
+    DEVICE_FORCEINLINE KVAddr decode_addr(
+        const AttentionParams<bf16>& p, const KVContext& c,
+        int batch, int kv_head, int kc, int d, bool valid, bool persist) {
+        if (p.new_k_ptr && valid && kc == kv_len(p, batch) - 1) {
+            KVAddr src = new_kv_addr(p, batch, kv_head, d);
+            if (persist) {
+                #pragma unroll
+                for (int j = 0; j < VEC; j++) {
+                    KVAddr value = new_kv_addr(p, batch, kv_head, d + j);
+                    store_new_kv(p, c, kc + 1, d + j, value);
+                }
+            }
+            return src;
+        }
+        int token = resolve_token(p, c, kc, valid);
+        return kv_addr_from_token(p, c, token, d);
+    }
+};

csrc/kernels/attn_mma_utils.cuh

@@ -99,22 +99,22 @@ __device__ __forceinline__ int swiz_col(int d, int r, int mask = 7) {
     return ((d >> 3) ^ (r & mask)) << 3 | (d & 7);
 }
 
-// cp.async: copy 16 bytes (8 bf16) from global to shared memory directly.
-__device__ __forceinline__ void cp_async_16(bf16* smem_ptr, const void* gmem_ptr) {
-    unsigned smem_addr = __cvta_generic_to_shared(smem_ptr);
-    asm volatile("cp.async.ca.shared.global [%0], [%1], 16;"
-                 :: "r"(smem_addr), "l"(gmem_ptr));
-}
-
 // Predicated cp.async: copy 16 bytes when `pred`, otherwise zero-fill.
-// src_size=0 → no bytes read from src, so out-of-bounds src address is safe.
+// BypassL1 defaults to .cg (L2 only); false selects .ca (L1 + L2).
+// src_size=0 means no bytes are read, so an out-of-bounds address is safe.
+template <bool BypassL1 = true>
 __device__ __forceinline__ void cp_async_16_pred(bf16* smem_ptr,
                                                   const void* gmem_ptr,
                                                   bool pred) {
     unsigned smem_addr = __cvta_generic_to_shared(smem_ptr);
     int src_size = pred ? 16 : 0;
-    asm volatile("cp.async.ca.shared.global [%0], [%1], 16, %2;"
-                 :: "r"(smem_addr), "l"(gmem_ptr), "r"(src_size));
+    if constexpr (BypassL1) {
+        asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;"
+                     :: "r"(smem_addr), "l"(gmem_ptr), "r"(src_size));
+    } else {
+        asm volatile("cp.async.ca.shared.global [%0], [%1], 16, %2;"
+                     :: "r"(smem_addr), "l"(gmem_ptr), "r"(src_size));
+    }
 }
 
 __device__ __forceinline__ void cp_async_commit() {
@@ -133,8 +133,8 @@ __device__ __forceinline__ void cp_async_wait_group() {
 // ---------------------------------------------------------------------------
 // Q-load: load query rows directly from global memory into mma A-operand
 // register layout. One call replaces ~15 duplicated lines in each MMA kernel.
-// stride_row is p.q_stride_h for decode (q_len=1, G heads) or
-//              p.q_stride_l for prefill (multi-q rows).
+// stride_row is p.q_h_stride for decode (q_len=1, G heads) or
+//              p.q_l_stride for prefill (multi-q rows).
 // ---------------------------------------------------------------------------
 template <int KD>
 __device__ inline void load_q_mma_frags(
@@ -198,9 +198,10 @@ __device__ inline void mma_softmax_tile(
     int kv0,
     int maxc0, int maxc1,
     int qrow0, int qrow1,
-    int mask_b_stride, int mask_h_stride, int mask_q_stride,
-    int mask_batch, int mask_head,
+    int mask_b_stride, int mask_h_stride, int mask_l_stride,
+    int mask_batch, int mask_head0, int mask_head1,
     const bool* __restrict__ mask,
+    bool valid0, bool valid1,
     float Sacc[Traits::NC8][4],
     float Oacc[Traits::DN8][4],
     float& m0, float& m1,
@@ -210,16 +211,16 @@ __device__ inline void mma_softmax_tile(
     int tid4 = lane & 3;
 
     float rmax0 = -FLT_MAX, rmax1 = -FLT_MAX;
-    int mask_base0 = mask_batch * mask_b_stride + mask_head * mask_h_stride + qrow0 * mask_q_stride;
-    int mask_base1 = mask_batch * mask_b_stride + mask_head * mask_h_stride + qrow1 * mask_q_stride;
+    int mask_base0 = mask_batch * mask_b_stride + mask_head0 * mask_h_stride + qrow0 * mask_l_stride;
+    int mask_base1 = mask_batch * mask_b_stride + mask_head1 * mask_h_stride + qrow1 * mask_l_stride;
     #pragma unroll
     for (int n8 = 0; n8 < Traits::NC8; n8++) {
         int cc = kv0 + n8 * 8 + 2 * tid4;
         int c1 = cc + 1;
-        bool b0 = (cc >= maxc0) || (HasMask && !mask[mask_base0 + cc]);
-        bool b1 = (c1 >= maxc0) || (HasMask && !mask[mask_base0 + c1]);
-        bool b2 = (cc >= maxc1) || (HasMask && !mask[mask_base1 + cc]);
-        bool b3 = (c1 >= maxc1) || (HasMask && !mask[mask_base1 + c1]);
+        bool b0 = !valid0 || (cc >= maxc0) || (HasMask && !mask[mask_base0 + cc]);
+        bool b1 = !valid0 || (c1 >= maxc0) || (HasMask && !mask[mask_base0 + c1]);
+        bool b2 = !valid1 || (cc >= maxc1) || (HasMask && !mask[mask_base1 + cc]);
+        bool b3 = !valid1 || (c1 >= maxc1) || (HasMask && !mask[mask_base1 + c1]);
         float s0 = b0 ? -FLT_MAX : Sacc[n8][0];
         float s1 = b1 ? -FLT_MAX : Sacc[n8][1];
         float s2 = b2 ? -FLT_MAX : Sacc[n8][2];

csrc/kernels/attn_paged_decode.cu

@@ -3,40 +3,84 @@
 
 torch::Tensor attn_paged_decode(
     torch::Tensor q,
-    torch::Tensor page_table,
     torch::Tensor k_cache,
     torch::Tensor v_cache,
-    int64_t page_size,
-    int64_t kv_len,
+    torch::Tensor req_to_token,
+    torch::Tensor req_pool_indices,
+    torch::Tensor kv_indptr,
+    c10::optional<torch::Tensor> new_k,
+    c10::optional<torch::Tensor> new_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,
+    c10::optional<torch::Tensor> out_buf
 ) {
-    PagedAttentionParams<bf16> p;
-    attn_pack_paged_params(q, page_table, k_cache, v_cache,
-                           page_size, kv_len, mask, causal_offset, scale, layout, p);
+    const at::cuda::OptionalCUDAGuard device_guard(device_of(q));
+    auto stream = at::cuda::getCurrentCUDAStream();
 
-    auto O = torch::empty_strided(q.sizes(), q.strides(), q.options());
-    auto O_view = (layout == 1) ? O.transpose(1, 2) : O;
-    p.o = (bf16*)O_view.data_ptr();
+    AttentionParams<bf16> p;
+    attn_pack_paged_decode_params(q, k_cache, v_cache,
+                                   req_to_token, req_pool_indices, kv_indptr,
+                                   new_k, new_v,
+                                   mask, causal_offset, scale, p);
 
-    alloc_split_partials(p);
-    DISPATCH_HEAD_DIM(p.head_dim, dispatch_paged_decode, p);
+    torch::Tensor O;
+    if (out_buf.has_value() && out_buf->defined()) {
+        TORCH_CHECK(out_buf->dtype() == q.dtype(), "out_buf dtype must match q");
+        TORCH_CHECK(out_buf->is_cuda() && out_buf->is_contiguous(),
+                    "out_buf must be a contiguous CUDA tensor");
+        TORCH_CHECK(out_buf->size(0) >= q.size(0), "out_buf batch too small");
+        TORCH_CHECK(out_buf->size(1) == q.size(1), "out_buf heads must match q");
+        TORCH_CHECK(out_buf->size(2) == q.size(2), "out_buf head_dim must match q");
+        TORCH_CHECK(q.is_contiguous(),
+                    "q must be contiguous when out_buf is provided");
+        O = out_buf.value().slice(0, 0, q.size(0));
+    } else {
+        O = torch::empty({q.size(0), q.size(1), q.size(2)}, q.options());
+    }
+    p.o_ptr = (bf16*)O.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;
+        int64_t ml_needed = (int64_t)p.batch * p.q_head * MAX_SPLITS * 2;
+        TORCH_CHECK(o_part_buf->numel() >= o_needed,
+                     "o_part_buf too small: need ", o_needed, " got ", o_part_buf->numel());
+        TORCH_CHECK(ml_part_buf->numel() >= ml_needed,
+                     "ml_part_buf too small: need ", ml_needed, " got ", ml_part_buf->numel());
+        TORCH_CHECK(o_part_buf->is_cuda() && ml_part_buf->is_cuda(),
+                    "split buffers must be CUDA tensors");
+        TORCH_CHECK(o_part_buf->is_contiguous() && ml_part_buf->is_contiguous(),
+                    "split buffers must be contiguous");
+        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_paged_decode, p, stream);
+    C10_CUDA_CHECK(cudaGetLastError());
     return O;
 }
 
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
     m.def("attn_paged_decode", &attn_paged_decode,
         py::arg("q"),
-        py::arg("page_table"),
         py::arg("k_cache"),
         py::arg("v_cache"),
-        py::arg("page_size"),
-        py::arg("kv_len"),
+        py::arg("req_to_token"),
+        py::arg("req_pool_indices"),
+        py::arg("kv_indptr"),
+        py::arg("new_k") = py::none(),
+        py::arg("new_v") = py::none(),
         py::arg("mask") = py::none(),
         py::arg("causal_offset") = -1,
         py::arg("scale") = 0.0,
-        py::arg("layout") = 0,
-        "Paged GQA decode — split-KV with direct page-table access.");
+        py::arg("o_part_buf") = py::none(),
+        py::arg("ml_part_buf") = py::none(),
+        py::arg("out_buf") = py::none(),
+        "SGLang-style paged decode: flat KV pool + req_to_token + kv_indptr.");
 }

csrc/kernels/attn_paged_decode_split_kv.cuh

@@ -1,153 +0,0 @@
-#pragma once
-#include <cuda_bf16.h>
-#include <float.h>
-#include "attn_common.h"
-#include "attn_warp_utils.cuh"
-constexpr int PDC_CHUNK = 64;
-
-template <int HEAD_DIM, bool IsCausal, bool HasMask>
-__global__ void paged_attn_decode_split_kv_kernel(PagedAttentionParams<bf16> p) {
-    int batch = blockIdx.x / p.kv_head;
-    int kv_head = blockIdx.x % p.kv_head;
-    int split = blockIdx.z;
-    int group_size = blockDim.y;
-    int q_head = kv_head * group_size + threadIdx.y;
-    int lane = threadIdx.x;
-    int hd_per_thread = p.head_dim / 32;
-
-    float q_reg[8];
-    int q_off = batch * p.q_stride_b + q_head * p.q_stride_h
-              + lane * hd_per_thread * p.q_stride_d;
-    #pragma unroll
-    for (int i = 0; i < hd_per_thread; i++)
-        q_reg[i] = __bfloat162float(p.q[q_off + i * p.q_stride_d]);
-
-    float m = -FLT_MAX, d = 0.0f, acc_reg[8] = {0.0f};
-
-    extern __shared__ __align__(16) bf16 k_smem[];
-
-    int chunks_total = (p.kv_len + PDC_CHUNK - 1) / PDC_CHUNK;
-    int chunks_per_split = (chunks_total + p.num_splits - 1) / p.num_splits;
-    int ch_begin = split * chunks_per_split;
-    int ch_end = min(chunks_total, ch_begin + chunks_per_split);
-
-    const int mask_base = batch * p.mask_b_stride + q_head * p.mask_h_stride;
-
-    for (int ci = ch_begin; ci < ch_end; ci++) {
-        int chunk_start = ci * PDC_CHUNK;
-        int this_chunk = min(PDC_CHUNK, p.kv_len - chunk_start);
-
-        int total = this_chunk * p.head_dim;
-        for (int i = threadIdx.y * 32 + lane; i < total;
-             i += blockDim.x * blockDim.y) {
-            int s = i / p.head_dim;
-            int d_dim = i % p.head_dim;
-            int pos = chunk_start + s;
-            int logical_page = pos / p.page_size;
-            int page_offset = pos % p.page_size;
-            int phys_page = p.page_table[batch * p.max_pages + logical_page];
-            if (phys_page >= 0) {
-                int64_t off = (int64_t)phys_page * p.page_size * p.kv_head * p.head_dim
-                            + (int64_t)page_offset * p.kv_head * p.head_dim
-                            + (int64_t)kv_head * p.head_dim
-                            + d_dim;
-                k_smem[i] = p.k_cache[off];
-            } else {
-                k_smem[i] = __float2bfloat16(0.0f);
-            }
-        }
-        __syncthreads();
-
-        for (int s = 0; s < this_chunk; s++) {
-            float partial = 0.0f;
-            #pragma unroll
-            for (int i = 0; i < hd_per_thread; i++)
-                partial += q_reg[i] * __bfloat162float(
-                    k_smem[s * p.head_dim + lane * hd_per_thread + i]);
-            partial = warp_reduce_sum(partial) * p.scale;
-
-            int kv_idx = chunk_start + s;
-            bool masked = false;
-            if constexpr (HasMask) {
-                if (!p.mask[mask_base + kv_idx])
-                    masked = true;
-            }
-            if constexpr (IsCausal) {
-                if (kv_idx > p.causal_offset)
-                    masked = true;
-            }
-            if (masked)
-                partial = -FLT_MAX;
-
-            float new_m = fmaxf(m, partial);
-            float alpha = __expf(m - new_m);
-            float beta  = __expf(partial - new_m);
-            d = d * alpha + beta;
-
-            int pos = chunk_start + s;
-            int logical_page = pos / p.page_size;
-            int page_offset = pos % p.page_size;
-            int phys_page = p.page_table[batch * p.max_pages + logical_page];
-            if (masked) {
-                #pragma unroll
-                for (int i = 0; i < hd_per_thread; i++)
-                    acc_reg[i] = fmaf(acc_reg[i], alpha, 0.0f);
-            } else if (phys_page >= 0) {
-                int64_t v_base = (int64_t)phys_page * p.page_size * p.kv_head * p.head_dim
-                               + (int64_t)page_offset * p.kv_head * p.head_dim
-                               + (int64_t)kv_head * p.head_dim;
-                #pragma unroll
-                for (int i = 0; i < hd_per_thread; i++)
-                    acc_reg[i] = fmaf(acc_reg[i], alpha,
-                                      __bfloat162float(p.v_cache[v_base + lane * hd_per_thread + i]) * beta);
-            } else {
-                #pragma unroll
-                for (int i = 0; i < hd_per_thread; i++)
-                    acc_reg[i] = fmaf(acc_reg[i], alpha, 0.0f);
-            }
-            m = new_m;
-        }
-        __syncthreads();
-    }
-
-    size_t bh = (size_t)batch * p.q_head + q_head;
-    size_t slot = bh * MAX_SPLITS + split;
-    int d0 = lane * hd_per_thread;
-    #pragma unroll
-    for (int i = 0; i < hd_per_thread; i++)
-        p.o_part[slot * p.head_dim + (d0 + i)] = acc_reg[i];
-    if (lane == 0) {
-        p.ml_part[slot * 2] = m;
-        p.ml_part[slot * 2 + 1] = d;
-    }
-}
-
-__global__ void paged_attn_decode_combine_kernel(PagedAttentionParams<bf16> p) {
-    int bh = blockIdx.x;
-    int d = threadIdx.x;
-    if (d >= p.head_dim) return;
-
-    int batch = bh / p.q_head;
-    int q_head = bh % p.q_head;
-
-    size_t split_base = (size_t)bh * 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 = batch * p.q_stride_b + q_head * p.q_stride_h + d * p.q_stride_d;
-    p.o[o_off] = __float2bfloat16(acc * inv);
-}

csrc/kernels/attn_paged_decode_split_kv_mma.cuh

@@ -1,182 +0,0 @@
-#pragma once
-#include <cfloat>
-#include <cuda_bf16.h>
-#include "attn_common.h"
-#include "attn_mma_utils.cuh"
-#include "attn_warp_utils.cuh"
-
-// Paged split-KV tensor-core decode via GQA head-packing.
-// Reads K/V directly from the page pool through a page table — one tile
-// (BC=32) fits within a single page (page_size >= 32), so the page-table
-// lookup happens once per tile for cp.async.
-//
-// IsCausal and HasMask are compile-time bools.
-template <typename Traits, bool IsCausal, bool HasMask>
-__global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams<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;
-
-    __shared__ __align__(16) bf16 sK[Traits::STAGES * Traits::BC * Traits::LD];
-    __shared__ __align__(16) bf16 sV[Traits::STAGES * Traits::BC * Traits::LD];
-
-    #pragma unroll
-    for (int i = lane; i < Traits::STAGES * Traits::BC * Traits::LD; i += 32) {
-        sK[i] = __float2bfloat16(0.0f);
-        sV[i] = __float2bfloat16(0.0f);
-    }
-    __syncwarp();
-
-    const int q_base = batch * p.q_stride_b + q_head0 * p.q_stride_h;
-    const int qra = gid;
-    const int qrb = gid + 8;
-    const bool va = qra < G, vb = qrb < G;
-    unsigned Qa[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 = (p.kv_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);
-
-    const int64_t page_stride = (int64_t)p.page_size * p.kv_head * Traits::HEAD_DIM;
-    const int64_t pos_stride  = (int64_t)p.kv_head * Traits::HEAD_DIM;
-    const int64_t head_off    = (int64_t)kv_head * Traits::HEAD_DIM;
-
-    // ---- Load tile lambda: paged addressing ----
-    // Unified per-element page-table lookup. When page_size >= BC, all
-    // elements in a tile share the same page, so the lookup is redundant
-    // but harmless (L1-cached). This avoids a branch on page_size.
-    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 < p.kv_len);
-            if constexpr (HasMask) {
-                valid = valid && p.mask[batch * p.mask_b_stride + kc];
-            }
-            int phys_page = valid ? p.page_table[batch * p.max_pages + kc] : 0;
-            valid = valid && (phys_page >= 0);
-            int page_off = kc % p.page_size;
-            int64_t gmem_base = (int64_t)phys_page * page_stride
-                              + (int64_t)page_off * pos_stride
-                              + head_off;
-            int off = r * Traits::LD + swiz_col(d, r, Traits::SWIZ_MASK);
-            cp_async_16_pred(&dK[off], &p.k_cache[gmem_base + d], valid);
-            cp_async_16_pred(&dV[off], &p.v_cache[gmem_base + d], valid);
-        }
-        cp_async_commit();
-    };
-
-    // ---- 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;
-
-        int maxc = IsCausal ? min(p.kv_len, p.causal_offset + 1) : p.kv_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);
-    }
-
-    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;
-        }
-    }
-}

csrc/kernels/attn_paged_prefill.cu

@@ -0,0 +1,51 @@
+#include "attn_dispatchers.cuh"
+#include "attn_entry_utils.cuh"
+
+torch::Tensor attn_paged_prefill(
+    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,
+    torch::Tensor q_tile_to_batch,
+    torch::Tensor q_tile_to_index,
+    c10::optional<torch::Tensor> mask,
+    int64_t causal_offset,
+    double scale
+) {
+    const at::cuda::OptionalCUDAGuard device_guard(device_of(q));
+    auto stream = at::cuda::getCurrentCUDAStream();
+
+    AttentionParams<bf16> p;
+    attn_pack_paged_prefill_params(q, k_cache, v_cache,
+                                     req_to_token, req_pool_indices,
+                                     kv_indptr, qo_indptr,
+                                     q_tile_to_batch, q_tile_to_index, mask,
+                                    causal_offset, scale, p);
+
+    auto O = torch::empty({q.size(0), q.size(1), q.size(2)}, q.options());
+    p.o_ptr = (bf16*)O.data_ptr();
+
+    DISPATCH_HEAD_DIM(p.head_dim, dispatch_paged_prefill, p, stream);
+    C10_CUDA_CHECK(cudaGetLastError());
+    return O;
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("attn_paged_prefill", &attn_paged_prefill,
+        py::arg("q"),
+        py::arg("k_cache"),
+        py::arg("v_cache"),
+        py::arg("req_to_token"),
+        py::arg("req_pool_indices"),
+        py::arg("kv_indptr"),
+        py::arg("qo_indptr"),
+        py::arg("q_tile_to_batch"),
+        py::arg("q_tile_to_index"),
+        py::arg("mask") = py::none(),
+        py::arg("causal_offset") = -1,
+        py::arg("scale") = 0.0,
+        "SGLang-style paged prefill: flat KV pool + ragged batch.");
+}

csrc/kernels/attn_prefill.cu

@@ -10,15 +10,19 @@ torch::Tensor attn_prefill(
     double scale,
     int64_t layout
 ) {
+    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.head_dim % 16 == 0, "head_dim must be multiple of 16");
 
     auto O = torch::empty_strided(q.sizes(), q.strides(), q.options());
-    auto O_view = (layout == 1) ? O.transpose(1, 2) : O;
-    p.o = (bf16*)O_view.data_ptr();
+    auto O_view = (layout == BLHD) ? O.transpose(1, 2) : O;
+    p.o_ptr = (bf16*)O_view.data_ptr();
 
-    DISPATCH_HEAD_DIM(p.head_dim, dispatch_prefill, p);
+    DISPATCH_HEAD_DIM(p.head_dim, dispatch_prefill, p, stream);
+    C10_CUDA_CHECK(cudaGetLastError());
     return O;
 }
 
@@ -30,6 +34,6 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
         py::arg("mask") = py::none(),
         py::arg("causal_offset") = -1,
         py::arg("scale") = 0.0,
-        py::arg("layout") = 0,
+        py::arg("layout") = (int64_t)BHLD,
         "GQA prefill (tensor-core mma on sm_80+, scalar fallback)");
 }

csrc/kernels/attn_prefill_split_q.cuh

@@ -2,13 +2,15 @@
 #include <cfloat>
 #include <cuda_bf16.h>
 #include "attn_common.h"
+#include "attn_layout_policies.cuh"
 
 using bf16 = __nv_bfloat16;
 
 // v9: group-split register blocking. G threads cooperate on one query row,
 // each owning HEAD_DIM/G dims of qreg[]/acc[]. IsCausal and HasMask are
 // compile-time bools — the compiler eliminates dead branches.
-// Templated on <HEAD_DIM, G, ROWS, P_BC, IsCausal, HasMask>.
+// Unified across contiguous and paged (SGLang flat-pool) K/V via KV.
+// Templated on <HEAD_DIM, KV, G, ROWS, P_BC, IsCausal, HasMask>.
 
 template <int G>
 __device__ __forceinline__ float group_reduce_sum(float v, unsigned mask) {
@@ -30,30 +32,37 @@ __device__ __forceinline__ void ld8(const bf16* p, float* o) {
     }
 }
 
-template <int HEAD_DIM, int G, int ROWS, int P_BC, bool IsCausal, bool HasMask>
+template <int HEAD_DIM, typename QSchedule, typename KV, int G, int ROWS, int P_BC,
+          bool IsCausal, bool HasMask>
 __global__ void attn_prefill_split_q_kernel_t(AttentionParams<bf16> p) {
     constexpr int DPT = HEAD_DIM / G;
 
-    int q_tile = blockIdx.x;
+    int batch, q_tile;
+    QSchedule::map_block(p, batch, q_tile);
+
     int q_head = blockIdx.y;
-    int batch  = blockIdx.z;
     int gpos   = threadIdx.x;  // 0..G-1  (which d-chunk)
     int row    = threadIdx.y;  // 0..ROWS-1
     int q_row  = q_tile * ROWS + row;
 
-    int kv_head = q_head / (p.q_head / p.kv_head);
+    // Per-request dims (from KV policy — paged reads kv_indptr/qo_indptr).
+    const int seq_len    = KV::kv_len(p, batch);
+    const int q_len      = QSchedule::q_len(p, batch);
+    const int causal_off = KV::causal_offset(p, batch, q_len);
+    const int kv_head = q_head / (p.q_head / p.kv_head);
+    const KVContext kctx = KV::template make_ctx<HEAD_DIM>(p, batch, kv_head);
 
     __shared__ __align__(16) bf16 sK[P_BC * HEAD_DIM];
     __shared__ __align__(16) bf16 sV[P_BC * HEAD_DIM];
 
     // Q: stride-based load [batch, q_head, q_len, head_dim]
+    const int q_base = QSchedule::q_base(p, batch, q_head);
     float qreg[DPT];
-    if (q_row < p.q_len) {
-        int q_off = batch * p.q_stride_b + q_head * p.q_stride_h
-                  + q_row * p.q_stride_l + gpos * DPT * p.q_stride_d;
+    if (q_row < q_len) {
+        int q_off = q_base + q_row * p.q_l_stride + gpos * DPT * p.q_d_stride;
 #pragma unroll
         for (int i = 0; i < DPT; i++)
-            qreg[i] = __bfloat162float(p.q[q_off + i * p.q_stride_d]);
+            qreg[i] = __bfloat162float(p.q_ptr[q_off + i * p.q_d_stride]);
     }
 
     float m = -FLT_MAX, l = 0.0f;
@@ -62,10 +71,8 @@ __global__ void attn_prefill_split_q_kernel_t(AttentionParams<bf16> p) {
     for (int i = 0; i < DPT; i++)
         acc[i] = 0.0f;
 
-    // KV: stride-based base
-    int kv_base = batch * p.kv_stride_b + kv_head * p.kv_stride_h;
     int mask_batch_base = batch * p.mask_b_stride + q_head * p.mask_h_stride;
-    int tiles   = (p.kv_len + P_BC - 1) / P_BC;
+    int tiles   = (seq_len + P_BC - 1) / P_BC;
     int tt      = G * ROWS;
     int lid     = row * G + gpos;
 
@@ -75,23 +82,25 @@ __global__ void attn_prefill_split_q_kernel_t(AttentionParams<bf16> p) {
 
     for (int ti = 0; ti < tiles; ti++) {
         int kv0  = ti * P_BC;
-        int tlen = min(P_BC, p.kv_len - kv0);
+        int tlen = min(P_BC, seq_len - kv0);
 
-        // Load K/V into shared memory from strided global
+        // Load K/V into shared memory (addressing via KV policy; paged
+        // guards empty slots with zero-fill).
         for (int i = lid; i < tlen * HEAD_DIM; i += tt) {
             int s = i / HEAD_DIM;
             int d_dim = i % HEAD_DIM;
-            int kv_idx = kv0 + s;
-            int g_off = kv_base + kv_idx * p.kv_stride_l + d_dim * p.kv_stride_d;
-            sK[i] = p.k[g_off];
-            sV[i] = p.v[g_off];
+            int kc = kv0 + s;
+            int token = KV::resolve_token(p, kctx, kc, true);
+            KVAddr a = KV::kv_addr_from_token(p, kctx, token, d_dim);
+            sK[i] = a.valid ? *reinterpret_cast<const bf16*>(a.k) : (bf16)0.f;
+            sV[i] = a.valid ? *reinterpret_cast<const bf16*>(a.v) : (bf16)0.f;
         }
         __syncthreads();
 
         int lim = tlen;
         if constexpr (IsCausal) {
-            if (q_row < p.q_len) {
-                int ep = q_row + p.causal_offset + 1;
+            if (q_row < q_len) {
+                int ep = causal_off + q_row + 1;
                 if (kv0 >= ep)
                     lim = 0;
                 else if (kv0 + tlen > ep)
@@ -99,7 +108,7 @@ __global__ void attn_prefill_split_q_kernel_t(AttentionParams<bf16> p) {
             }
         }
 
-        int mask_row_base = mask_batch_base + q_row * p.mask_q_stride;
+        int mask_row_base = mask_batch_base + q_row * p.mask_l_stride;
         for (int s = 0; s < lim; s++) {
             const bf16* kr = sK + s * HEAD_DIM + gpos * DPT;
             float part = 0.0f;
@@ -138,12 +147,11 @@ __global__ void attn_prefill_split_q_kernel_t(AttentionParams<bf16> p) {
         __syncthreads();
     }
 
-    if (q_row < p.q_len) {
-        int o_off = batch * p.q_stride_b + q_head * p.q_stride_h
-                  + q_row * p.q_stride_l + gpos * DPT * p.q_stride_d;
+    if (q_row < q_len) {
+        int o_off = q_base + q_row * p.q_l_stride + gpos * DPT * p.q_d_stride;
         float rl = (l > 1e-20f) ? (1.0f / l) : 0.0f;
 #pragma unroll
         for (int i = 0; i < DPT; i++)
-            p.o[o_off + i * p.q_stride_d] = __float2bfloat16(acc[i] * rl);
+            p.o_ptr[o_off + i * p.q_d_stride] = __float2bfloat16(acc[i] * rl);
     }
 }

csrc/kernels/attn_prefill_split_q_mma.cuh

@@ -2,17 +2,21 @@
 #include <cfloat>
 #include <cuda_bf16.h>
 #include "attn_common.h"
+#include "attn_layout_policies.cuh"
 #include "attn_mma_utils.cuh"
 
-// Tensor-core prefill flash attention (raw mma.sync PTX).
+// Tensor-core prefill flash attention (raw mma.sync PTX), unified across
+// contiguous and paged (SGLang flat-pool) K/V via the KV template parameter.
 // One warp owns BR=16 query rows. S = Q@K^T and O = P@V run on bf16 tensor
 // cores via mma.sync.m16n8k16 (f32 accumulate).
 //
+// KV = ContigKV (dense [batch, kv_head, kv_len, head_dim]) or PagedKV
+//      (flat pool + req_to_token, ragged batches via qo_indptr/kv_indptr).
 // IsCausal and HasMask are compile-time bools — the compiler eliminates all
 // dead branches in the inner compute loop (FA2-style).
 //
 // Traits = KernelTraits<HEAD_DIM, BC, WARPS=4, STAGES=2>.
-template <typename Traits, bool IsCausal, bool HasMask>
+template <typename Traits, typename QSchedule, typename KV, bool IsCausal, bool HasMask>
 __global__ void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
     const int warp = threadIdx.x / 32;
     const int lane = threadIdx.x % 32;
@@ -20,9 +24,16 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
     const int tid4 = lane & 3;   // 0..3
 
     const int q_head = blockIdx.y;
-    const int batch = blockIdx.z;
+    int batch, q_tile;
+    QSchedule::map_block(p, batch, q_tile);
     const int kv_head = q_head / (p.q_head / p.kv_head);
-    const int qrow0 = (blockIdx.x * Traits::WARPS + warp) * Traits::BR;
+    const int qrow0 = (q_tile * Traits::WARPS + warp) * Traits::BR;
+
+    // Per-request dims (from KV policy — paged reads kv_indptr/qo_indptr).
+    const int seq_len    = KV::kv_len(p, batch);
+    const int q_len      = QSchedule::q_len(p, batch);
+    const int causal_off = KV::causal_offset(p, batch, q_len);
+    const KVContext kctx = KV::template make_ctx<Traits::HEAD_DIM>(p, batch, kv_head);
 
     // Static shared memory: double-buffered K/V (no sQ — Q goes direct
     // to registers in mma A-operand layout).
@@ -30,12 +41,12 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
     __shared__ __align__(16) bf16 sV[Traits::STAGES * Traits::BC * Traits::LD];
 
     // Load Q fragments straight from global into mma A-operand layout.
-    const int q_base = batch * p.q_stride_b + q_head * p.q_stride_h;
+    const int q_base = QSchedule::q_base(p, batch, q_head);
     const int qra = qrow0 + gid;
     const int qrb = qrow0 + gid + 8;
-    const bool va = qra < p.q_len, vb = qrb < p.q_len;
+    const bool va = qra < q_len, vb = qrb < q_len;
     unsigned Qa[Traits::KD][4];
-    load_q_mma_frags<Traits::KD>(p.q + q_base, p.q_stride_l, p.q_stride_d,
+    load_q_mma_frags<Traits::KD>(p.q_ptr + q_base, p.q_l_stride, p.q_d_stride,
                                   qra, qrb, va, vb, tid4, Qa);
 
     float Oacc[Traits::DN8][4];
@@ -44,17 +55,15 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
         Oacc[j][0] = Oacc[j][1] = Oacc[j][2] = Oacc[j][3] = 0.0f;
     float m0 = -FLT_MAX, m1 = -FLT_MAX, l0 = 0.0f, l1 = 0.0f;
 
-    // KV: stride-based base
-    const int kv_base = batch * p.kv_stride_b + kv_head * p.kv_stride_h;
-    const int tiles = (p.kv_len + Traits::BC - 1) / Traits::BC;
+    const int tiles = (seq_len + Traits::BC - 1) / Traits::BC;
     const int qr0 = qrow0 + gid;
     const int qr1 = qrow0 + gid + 8;
 
     // Causal tile-skip bounds (dead code when IsCausal == false)
-    const int max_kv = qrow0 + Traits::BR - 1 + p.causal_offset;
+    const int max_kv = qrow0 + Traits::BR - 1 + causal_off;
     const int block_max_kv =
-        blockIdx.x * Traits::WARPS * Traits::BR + Traits::WARPS * Traits::BR - 1
-        + p.causal_offset;
+        q_tile * Traits::WARPS * Traits::BR + Traits::WARPS * Traits::BR - 1
+        + causal_off;
 
     int t_end = tiles - 1;
     if constexpr (IsCausal) {
@@ -62,7 +71,7 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
         if (bt < t_end) t_end = bt;
     }
 
-    // ---- Load tile lambda: predicated cp.async ----
+    // ---- 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;
@@ -72,11 +81,12 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
              i += Traits::NUM_THREADS * Traits::VEC) {
             int r = i / Traits::HEAD_DIM, d = i % Traits::HEAD_DIM;
             int kc = kv0 + r;
-            bool valid = kc < p.kv_len;
+            bool valid = kc < seq_len;
+            int token = KV::resolve_token(p, kctx, kc, valid);
+            KVAddr a = KV::kv_addr_from_token(p, kctx, token, d);
             int off = r * Traits::LD + swiz_col(d, r, Traits::SWIZ_MASK);
-            int g_off = kv_base + kc * p.kv_stride_l + d * p.kv_stride_d;
-            cp_async_16_pred(&dK[off], &p.k[g_off], valid);
-            cp_async_16_pred(&dV[off], &p.v[g_off], valid);
+            cp_async_16_pred(&dK[off], a.k, a.valid);
+            cp_async_16_pred(&dV[off], a.v, a.valid);
         }
         cp_async_commit();
     };
@@ -108,15 +118,16 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
                 Sacc[n8][0] *= p.scale, Sacc[n8][1] *= p.scale,
                 Sacc[n8][2] *= p.scale, Sacc[n8][3] *= p.scale;
 
-            int maxc0 = IsCausal ? min(p.kv_len, qr0 + p.causal_offset + 1)
-                                 : p.kv_len;
-            int maxc1 = IsCausal ? min(p.kv_len, qr1 + p.causal_offset + 1)
-                                 : p.kv_len;
+            int maxc0 = IsCausal ? min(seq_len, causal_off + qr0 + 1)
+                                 : seq_len;
+            int maxc1 = IsCausal ? min(seq_len, causal_off + qr1 + 1)
+                                 : seq_len;
             mma_softmax_tile<Traits, HasMask>(kv0, maxc0, maxc1,
                                                qr0, qr1,
-                                               p.mask_b_stride, p.mask_h_stride, p.mask_q_stride,
-                                               batch, q_head,
+                                               p.mask_b_stride, p.mask_h_stride, p.mask_l_stride,
+                                               batch, q_head, q_head,
                                                p.mask,
+                                               va, vb,
                                                Sacc, Oacc, m0, m1, l0, l1, lane);
 
             mma_pv_accumulate<Traits>(Sacc, bV, lane, Oacc);
@@ -126,21 +137,21 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
     // ---- write output: packed bf16x2 stores ----
     float rl0 = (l0 > 1e-20f) ? (1.0f / l0) : 0.0f;
     float rl1 = (l1 > 1e-20f) ? (1.0f / l1) : 0.0f;
-    const int o_base = batch * p.q_stride_b + q_head * p.q_stride_h;
+    const int o_base = QSchedule::q_base(p, batch, q_head);
     #pragma unroll
     for (int dn8 = 0; dn8 < Traits::DN8; dn8++) {
         int d = dn8 * 8 + 2 * tid4;
-        if (qr0 < p.q_len) {
+        if (qr0 < q_len) {
             __nv_bfloat162 v = __floats2bfloat162_rn(Oacc[dn8][0] * rl0,
                                                       Oacc[dn8][1] * rl0);
             *reinterpret_cast<__nv_bfloat162*>(
-                &p.o[o_base + qr0 * p.q_stride_l + d * p.q_stride_d]) = v;
+                &p.o_ptr[o_base + qr0 * p.q_l_stride + d * p.q_d_stride]) = v;
         }
-        if (qr1 < p.q_len) {
+        if (qr1 < q_len) {
             __nv_bfloat162 v = __floats2bfloat162_rn(Oacc[dn8][2] * rl1,
                                                       Oacc[dn8][3] * rl1);
             *reinterpret_cast<__nv_bfloat162*>(
-                &p.o[o_base + qr1 * p.q_stride_l + d * p.q_stride_d]) = v;
+                &p.o_ptr[o_base + qr1 * p.q_l_stride + d * p.q_d_stride]) = v;
         }
     }
 }

csrc/kernels/fp8_mm.cu

@@ -0,0 +1,817 @@
+// Fused BF16 -> E4M3 MMA -> BF16 matrix multiplication
+
+#include <torch/extension.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <cuda_fp8.h>
+#include <cuda_runtime.h>
+#include <cstdint>
+#include <mutex>
+#include <unordered_map>
+
+namespace {
+
+constexpr int kMmaK = 32;
+constexpr int kWarps = 8;
+// Fast forward path: 128x64 CTA, 64x16 warp tile, 2-stage pipeline, dynamic
+// shared memory. Mirrors the CUTLASS 58_ada_fp8_gemm threadblock geometry
+// while keeping the fused BF16->FP8 quantize path. The FP8 tile overwrites
+// the BF16 staging area in place. L20 opts in to only 101376 B shared per
+// block; K=32 keeps the footprint at 24576 B so four CTAs/SM stay resident.
+constexpr int kFastBlockM = 128;
+constexpr int kFastBlockN = 64;
+constexpr int kFastK = 32;
+constexpr int kFastStages = 2;
+constexpr int kFastSmemBytes =
+    kFastStages * (kFastBlockM * kFastK * 2 + kFastBlockN * kFastK * 2);
+
+__device__ __forceinline__ unsigned pack_fp8x4_vector(float x0, float x1,
+                                                       float x2, float x3) {
+    const auto low = __nv_cvt_float2_to_fp8x2(
+        make_float2(x0, x1), __NV_SATFINITE, __NV_E4M3);
+    const auto high = __nv_cvt_float2_to_fp8x2(
+        make_float2(x2, x3), __NV_SATFINITE, __NV_E4M3);
+    return static_cast<unsigned>(low) | (static_cast<unsigned>(high) << 16);
+}
+
+
+__device__ __forceinline__ void mma_fp8_16832(float d[4],
+                                               const unsigned a[4],
+                                               const unsigned b[2]) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 890
+    asm volatile(
+        "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
+        "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};"
+        : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3])
+        : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]),
+          "r"(b[0]), "r"(b[1]));
+#endif
+}
+
+__device__ __forceinline__ void atomic_max_float(float* destination,
+                                                  float value) {
+    if (destination)
+        atomicMax(reinterpret_cast<unsigned*>(destination), __float_as_uint(value));
+}
+
+__device__ __forceinline__ float warp_reduce_max(float value) {
+#pragma unroll
+    for (int offset = 16; offset; offset >>= 1) {
+        value = fmaxf(value, __shfl_xor_sync(0xffffffffu, value, offset));
+    }
+    return value;
+}
+
+// Block-wide max reduction of a per-warp tracked value, then an atomic
+// update of the global amax slot when `track` is set.
+template <int NWarps>
+__device__ __forceinline__ void block_reduce_amax(float& local, float* slots,
+                                                  int warp, int lane,
+                                                  bool track, float* global) {
+    local = warp_reduce_max(local);
+    if (lane == 0) slots[warp] = local;
+    __syncthreads();
+    if (warp == 0) {
+        float value = lane < NWarps ? slots[lane] : 0.0f;
+        value = warp_reduce_max(value);
+        if (lane == 0 && track && global) atomic_max_float(global, value);
+    }
+}
+
+// One thread moves eight BF16 values (16 bytes). The async copy is issued
+// through a uint4-shaped pointer so the source and destination are both
+// naturally 128-bit aligned for contiguous forward GEMMs.
+__device__ __forceinline__ void cp_async_bf16_8(
+    __nv_bfloat16* destination, const __nv_bfloat16* source, bool valid) {
+    const unsigned shared_address = __cvta_generic_to_shared(destination);
+    const uint4* source_vec = reinterpret_cast<const uint4*>(source);
+    asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;"
+                 :: "r"(shared_address), "l"(source_vec),
+                    "r"(valid ? 16 : 0));
+}
+
+template <bool TrackAmax = true>
+__device__ __forceinline__ unsigned load_fp8x4_from_bf16(
+    const __nv_bfloat16* source, float scale_inv, float& amax,
+    bool track_amax = true) {
+    float x0 = __bfloat162float(source[0]);
+    float x1 = __bfloat162float(source[1]);
+    float x2 = __bfloat162float(source[2]);
+    float x3 = __bfloat162float(source[3]);
+    if constexpr (TrackAmax) {
+        if (track_amax) {
+        amax = fmaxf(amax, fmaxf(fabsf(x0), fmaxf(fabsf(x1),
+                             fmaxf(fabsf(x2), fabsf(x3)))));
+        }
+    }
+    return pack_fp8x4_vector(x0 * scale_inv, x1 * scale_inv,
+                             x2 * scale_inv, x3 * scale_inv);
+}
+
+
+template <bool AddBias, bool TrackAmax>
+__global__ void fused_fp8_gemm_fast_kernel(
+    const __nv_bfloat16* __restrict__ a,
+    const __nv_bfloat16* __restrict__ b,
+    __nv_bfloat16* __restrict__ out,
+    const __nv_bfloat16* __restrict__ bias,
+    const float* __restrict__ scale_a,
+    const float* __restrict__ scale_b,
+    float* __restrict__ amax_a,
+    float* __restrict__ amax_b,
+    int64_t m, int64_t n, int64_t k) {
+    extern __shared__ char smem[];
+    constexpr int a_stride = kFastBlockM * kFastK;
+    constexpr int b_stride = kFastBlockN * kFastK;
+    constexpr int b_bf16_offset = kFastStages * a_stride;
+    auto* a_bf16 = reinterpret_cast<__nv_bfloat16*>(smem);
+    auto* b_bf16 =
+        reinterpret_cast<__nv_bfloat16*>(smem + b_bf16_offset * 2);
+    __shared__ float warp_amax_a[kWarps];
+    __shared__ float warp_amax_b[kWarps];
+
+    const int tid = threadIdx.x;
+    const int warp = tid >> 5;
+    const int lane = tid & 31;
+    const int group = lane >> 2;
+    const int thread_in_group = lane & 3;
+    constexpr int warps_n = kFastBlockN / 16;
+    const int warp_m = warp / warps_n;
+    const int warp_n = warp % warps_n;
+    const int64_t row_base =
+        blockIdx.y * kFastBlockM + warp_m * 64 + group;
+    const int64_t output_col =
+        blockIdx.x * kFastBlockN + warp_n * 16 + thread_in_group * 2;
+    const float sa = *scale_a;
+    const float sb = *scale_b;
+    const float inv_a = 1.0f / sa;
+    const float inv_b = 1.0f / sb;
+    float local_amax_a = 0.0f;
+    float local_amax_b = 0.0f;
+    float acc[4 * 4 * 2] = {};
+
+    const bool track_amax_a = TrackAmax && blockIdx.x == 0;
+    const bool track_amax_b = TrackAmax && blockIdx.y == 0;
+    // Each thread issues 8 A chunks and 4 B chunks of 8 BF16 (16B) per stage.
+    auto load_tile = [&](int stage, int64_t k_base) {
+        const int r0 = tid >> 2;
+        const int c0 = (tid & 3) * 8;
+#pragma unroll
+        for (int j = 0; j < kFastK / 32; ++j) {
+            const int col = c0 + 32 * j;
+            const bool full_chunk = k_base + col + 7 < k;
+            const int64_t a_row = blockIdx.y * kFastBlockM + r0;
+            const int64_t b_row = blockIdx.x * kFastBlockN + r0;
+            auto* a_dst = &a_bf16[stage * a_stride + r0 * kFastK + col];
+            auto* b_dst = &b_bf16[stage * b_stride + r0 * kFastK + col];
+            const auto* a_ptr = a + a_row * k + k_base + col;
+            const auto* b_ptr = b + b_row * k + k_base + col;
+            const bool full_a = a_row < m && full_chunk;
+            const bool full_b = b_row < n && full_chunk;
+            const bool aligned_a =
+                (reinterpret_cast<uintptr_t>(a_ptr) & 15) == 0;
+            const bool aligned_b =
+                (reinterpret_cast<uintptr_t>(b_ptr) & 15) == 0;
+            if (full_a && aligned_a) {
+                cp_async_bf16_8(a_dst, a_ptr, true);
+            } else {
+#pragma unroll
+                for (int i = 0; i < 8; ++i) {
+                    a_dst[i] = a_row < m && k_base + col + i < k
+                                   ? a_ptr[i]
+                                   : __float2bfloat16(0.0f);
+                }
+            }
+            if (full_b && aligned_b) {
+                cp_async_bf16_8(b_dst, b_ptr, true);
+            } else {
+#pragma unroll
+                for (int i = 0; i < 8; ++i) {
+                    b_dst[i] = b_row < n && k_base + col + i < k
+                                   ? b_ptr[i]
+                                   : __float2bfloat16(0.0f);
+                }
+            }
+            if (r0 + 64 < kFastBlockM) {
+                const int64_t a_row_hi = blockIdx.y * kFastBlockM + r0 + 64;
+                auto* a_dst_hi =
+                    &a_bf16[stage * a_stride + (r0 + 64) * kFastK + col];
+                const auto* a_ptr_hi = a + a_row_hi * k + k_base + col;
+                const bool full_a_hi = a_row_hi < m && full_chunk;
+                const bool aligned_a_hi =
+                    (reinterpret_cast<uintptr_t>(a_ptr_hi) & 15) == 0;
+                if (full_a_hi && aligned_a_hi) {
+                    cp_async_bf16_8(a_dst_hi, a_ptr_hi, true);
+                } else {
+#pragma unroll
+                    for (int i = 0; i < 8; ++i) {
+                        a_dst_hi[i] = a_row_hi < m && k_base + col + i < k
+                                          ? a_ptr_hi[i]
+                                          : __float2bfloat16(0.0f);
+                    }
+                }
+            }
+        }
+    };
+
+    // Quantize must place each 4-BP8 group at the byte offset the MMA
+    // fragment reads: 8*(lane&3) + 64*k_seg for a row. With in-place storage
+    // (fp8 element k lives at byte 2k), the BF16 column of a group is
+    // 4*(tid&7) + 32*j, so partition by 4-element groups instead of the
+    // 8-element cp.async chunks.
+    auto quantize_tile = [&](int stage) {
+        const int r0 = tid >> 3;
+        const int c0 = (tid & 7) * 4;
+#pragma unroll
+        for (int s = 0; s < 4; ++s) {
+            const int row = r0 + 32 * s;
+            auto* a_src = &a_bf16[stage * a_stride + row * kFastK + c0];
+            auto* a_dst = reinterpret_cast<unsigned*>(a_src);
+#pragma unroll
+            for (int j = 0; j < kFastK / 32; ++j) {
+                a_dst[16 * j] = load_fp8x4_from_bf16<TrackAmax>(
+                    a_src + 32 * j, inv_a, local_amax_a, track_amax_a);
+            }
+        }
+#pragma unroll
+        for (int s = 0; s < 2; ++s) {
+            const int row = r0 + 32 * s;
+            auto* b_src = &b_bf16[stage * b_stride + row * kFastK + c0];
+            auto* b_dst = reinterpret_cast<unsigned*>(b_src);
+#pragma unroll
+            for (int j = 0; j < kFastK / 32; ++j) {
+                b_dst[16 * j] = load_fp8x4_from_bf16<TrackAmax>(
+                    b_src + 32 * j, inv_b, local_amax_b, track_amax_b);
+            }
+        }
+    };
+
+    const int64_t tile_count = (k + kFastK - 1) / kFastK;
+    load_tile(0, 0);
+    asm volatile("cp.async.commit_group;");
+    if (tile_count > 1) {
+        load_tile(1, kFastK);
+        asm volatile("cp.async.commit_group;");
+    }
+    for (int64_t tile_index = 0; tile_index < tile_count; ++tile_index) {
+        const int stage = static_cast<int>(tile_index % kFastStages);
+        if (tile_index + 1 == tile_count) {
+            asm volatile("cp.async.wait_group 0;");
+        } else {
+            asm volatile("cp.async.wait_group 1;");
+        }
+        // wait_group only waits for this thread's async copies. All threads
+        // must finish loading before the tile is read by the CTA.
+        __syncthreads();
+        quantize_tile(stage);
+        __syncthreads();
+
+        // Four m16n8k32 MMA segments per 128-K stage.
+#pragma unroll
+        for (int k_seg = 0; k_seg < kFastK / kMmaK; ++k_seg) {
+            const int frag_col = thread_in_group * 4 + k_seg * 32;
+#pragma unroll
+            for (int nt = 0; nt < 2; ++nt) {
+                const int b_row = warp_n * 16 + nt * 8 + group;
+                unsigned b_frag[2];
+                b_frag[0] = *reinterpret_cast<const unsigned*>(
+                    &b_bf16[stage * b_stride + b_row * kFastK + frag_col]);
+                b_frag[1] = *reinterpret_cast<const unsigned*>(
+                    &b_bf16[stage * b_stride + b_row * kFastK + frag_col + 16]);
+#pragma unroll
+                for (int mt = 0; mt < 4; ++mt) {
+                    const int a_row0 = warp_m * 64 + mt * 16 + group;
+                    unsigned a_frag[4];
+                    a_frag[0] = *reinterpret_cast<const unsigned*>(
+                        &a_bf16[stage * a_stride + a_row0 * kFastK + frag_col]);
+                    a_frag[1] = *reinterpret_cast<const unsigned*>(
+                        &a_bf16[stage * a_stride + (a_row0 + 8) * kFastK + frag_col]);
+                    a_frag[2] = *reinterpret_cast<const unsigned*>(
+                        &a_bf16[stage * a_stride + a_row0 * kFastK + frag_col + 16]);
+                    a_frag[3] = *reinterpret_cast<const unsigned*>(
+                        &a_bf16[stage * a_stride + (a_row0 + 8) * kFastK + frag_col + 16]);
+                    mma_fp8_16832(acc + (nt * 4 + mt) * 4, a_frag, b_frag);
+                }
+            }
+        }
+        __syncthreads();
+        if (tile_index + 2 < tile_count) {
+            load_tile(stage, (tile_index + 2) * kFastK);
+            asm volatile("cp.async.commit_group;");
+        }
+    }
+
+    if constexpr (TrackAmax) {
+    block_reduce_amax<kWarps>(local_amax_a, warp_amax_a, warp, lane,
+                              track_amax_a, amax_a);
+    block_reduce_amax<kWarps>(local_amax_b, warp_amax_b, warp, lane,
+                              track_amax_b, amax_b);
+    }
+
+    const float output_scale = sa * sb;
+#pragma unroll
+    for (int nt = 0; nt < 2; ++nt) {
+        const int64_t col = output_col + nt * 8;
+#pragma unroll
+        for (int mt = 0; mt < 4; ++mt) {
+            const int64_t row0 = row_base + mt * 16;
+            const int64_t row1 = row0 + 8;
+            float* tile_acc = acc + (nt * 4 + mt) * 4;
+            if (col < n) {
+                float bias0 = 0.0f;
+                float bias1 = 0.0f;
+                if constexpr (AddBias) {
+                    bias0 = __bfloat162float(bias[col]);
+                    if (col + 1 < n)
+                        bias1 = __bfloat162float(bias[col + 1]);
+                }
+                if (row0 < m) {
+                    out[row0 * n + col] =
+                        __float2bfloat16(tile_acc[0] * output_scale + bias0);
+                    if (col + 1 < n)
+                        out[row0 * n + col + 1] = __float2bfloat16(
+                            tile_acc[1] * output_scale + bias1);
+                }
+                if (row1 < m) {
+                    out[row1 * n + col] =
+                        __float2bfloat16(tile_acc[2] * output_scale + bias0);
+                    if (col + 1 < n)
+                        out[row1 * n + col + 1] = __float2bfloat16(
+                            tile_acc[3] * output_scale + bias1);
+                }
+            }
+        }
+    }
+}
+
+// Pre-quantized FP8-in path: FP8 A/B read straight into shared memory (no
+// BF16 staging, no inline quantization), FP32 accumulation, BF16 output.
+// Same 128x64 CTA / 64x16 warp tile geometry as the fused kernel; the fp8
+// tile is compact (row = kFastK bytes) so MMA fragments read directly.
+constexpr int kPqBlockM = 128;
+constexpr int kPqBlockN = 64;
+constexpr int kPqK = 32;
+constexpr int kPqStages = 3;
+
+template <typename T>
+__device__ __forceinline__ void cp_async_16b(T* destination,
+                                             const T* source, bool valid) {
+    const unsigned shared_address = __cvta_generic_to_shared(destination);
+    const uint4* source_vec = reinterpret_cast<const uint4*>(source);
+    asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;"
+                 :: "r"(shared_address), "l"(source_vec),
+                    "r"(valid ? 16 : 0));
+}
+
+template <bool OutFp8>
+__global__ void fp8_mm_pq_kernel(
+    const __nv_fp8_e4m3* __restrict__ a,
+    const __nv_fp8_e4m3* __restrict__ b,
+    __nv_bfloat16* __restrict__ out_bf16,
+    __nv_fp8_e4m3* __restrict__ out_fp8,
+    const float scale, const float out_scale,
+    int64_t m, int64_t n, int64_t k) {
+    __shared__ __align__(16) __nv_fp8_e4m3 a_tile[kPqStages][kPqBlockM][kPqK];
+    __shared__ __align__(16) __nv_fp8_e4m3 b_tile[kPqStages][kPqBlockN][kPqK];
+
+    const int tid = threadIdx.x;
+    const int warp = tid >> 5;
+    const int lane = tid & 31;
+    const int group = lane >> 2;
+    const int thread_in_group = lane & 3;
+    constexpr int warps_n = kPqBlockN / 16;
+    const int warp_m = warp / warps_n;
+    const int warp_n = warp % warps_n;
+    const int64_t row_base = blockIdx.y * kPqBlockM + warp_m * 64 + group;
+    const int64_t output_col =
+        blockIdx.x * kPqBlockN + warp_n * 16 + thread_in_group * 2;
+    float acc[4 * 4 * 2] = {};
+
+    // One A chunk (16 FP8) per thread covers the 128x32 tile; the first 128
+    // threads issue the 64x32 B chunks.
+    auto load_tile = [&](int stage, int64_t k_base) {
+        const int r0 = tid >> 1;
+        const int c0 = (tid & 1) * 16;
+        const bool full_chunk = k_base + c0 + 15 < k;
+        const int64_t a_row = blockIdx.y * kPqBlockM + r0;
+        auto* a_dst = &a_tile[stage][r0][c0];
+        const auto* a_ptr = a + a_row * k + k_base + c0;
+        const bool full_a = a_row < m && full_chunk;
+        const bool aligned_a =
+            (reinterpret_cast<uintptr_t>(a_ptr) & 15) == 0;
+        if (full_a && aligned_a) {
+            cp_async_16b(a_dst, a_ptr, true);
+        } else {
+#pragma unroll
+            for (int i = 0; i < 16; ++i) {
+                a_dst[i] = a_row < m && k_base + c0 + i < k
+                               ? a_ptr[i]
+                               : __nv_fp8_e4m3(0.0f);
+            }
+        }
+        if (tid < 128) {
+            const int64_t b_row = blockIdx.x * kPqBlockN + r0;
+            auto* b_dst = &b_tile[stage][r0][c0];
+            const auto* b_ptr = b + b_row * k + k_base + c0;
+            const bool full_b = b_row < n && full_chunk;
+            const bool aligned_b =
+                (reinterpret_cast<uintptr_t>(b_ptr) & 15) == 0;
+            if (full_b && aligned_b) {
+                cp_async_16b(b_dst, b_ptr, true);
+            } else {
+#pragma unroll
+                for (int i = 0; i < 16; ++i) {
+                    b_dst[i] = b_row < n && k_base + c0 + i < k
+                                   ? b_ptr[i]
+                                   : __nv_fp8_e4m3(0.0f);
+                }
+            }
+        }
+    };
+
+    const int64_t tile_count = (k + kPqK - 1) / kPqK;
+    load_tile(0, 0);
+    asm volatile("cp.async.commit_group;");
+    if (tile_count > 1) {
+        load_tile(1, kPqK);
+        asm volatile("cp.async.commit_group;");
+    }
+    if (tile_count > 2) {
+        load_tile(2, 2 * kPqK);
+        asm volatile("cp.async.commit_group;");
+    }
+    for (int64_t tile_index = 0; tile_index < tile_count; ++tile_index) {
+        const int stage = static_cast<int>(tile_index % kPqStages);
+        const int64_t remaining = tile_count - tile_index - 1;
+        if (remaining >= 2) {
+            asm volatile("cp.async.wait_group 2;");
+        } else if (remaining == 1) {
+            asm volatile("cp.async.wait_group 1;");
+        } else {
+            asm volatile("cp.async.wait_group 0;");
+        }
+        // Barrier 1: every thread's cp.async for this stage is complete
+        // before any thread reads tiles written by other threads.
+        __syncthreads();
+
+#pragma unroll
+        for (int k_seg = 0; k_seg < kPqK / kMmaK; ++k_seg) {
+            const int frag_col = thread_in_group * 4 + k_seg * 32;
+#pragma unroll
+            for (int nt = 0; nt < 2; ++nt) {
+                const int b_row = warp_n * 16 + nt * 8 + group;
+                unsigned b_frag[2];
+                b_frag[0] = *reinterpret_cast<const unsigned*>(
+                    &b_tile[stage][b_row][frag_col]);
+                b_frag[1] = *reinterpret_cast<const unsigned*>(
+                    &b_tile[stage][b_row][frag_col + 16]);
+#pragma unroll
+                for (int mt = 0; mt < 4; ++mt) {
+                    const int a_row0 = warp_m * 64 + mt * 16 + group;
+                    unsigned a_frag[4];
+                    a_frag[0] = *reinterpret_cast<const unsigned*>(
+                        &a_tile[stage][a_row0][frag_col]);
+                    a_frag[1] = *reinterpret_cast<const unsigned*>(
+                        &a_tile[stage][a_row0 + 8][frag_col]);
+                    a_frag[2] = *reinterpret_cast<const unsigned*>(
+                        &a_tile[stage][a_row0][frag_col + 16]);
+                    a_frag[3] = *reinterpret_cast<const unsigned*>(
+                        &a_tile[stage][a_row0 + 8][frag_col + 16]);
+                    mma_fp8_16832(acc + (nt * 4 + mt) * 4, a_frag, b_frag);
+                }
+            }
+        }
+        // Barrier 2: every thread finished reading this stage's tiles before
+        // the prefetch for the (i+3)-th tile overwrites them.
+        __syncthreads();
+        if (tile_index + 3 < tile_count) {
+            load_tile(stage, (tile_index + 3) * kPqK);
+            asm volatile("cp.async.commit_group;");
+        }
+    }
+
+    const float output_scale = scale * out_scale;
+#pragma unroll
+    for (int nt = 0; nt < 2; ++nt) {
+        const int64_t col = output_col + nt * 8;
+        // Per-row store: FP8 packs two adjacent columns into one 16-bit
+        // write; the BF16 path writes two scalars. Boundary columns fall
+        // back to a scalar convert so the pack never crosses the row edge.
+        auto store_out = [&](int64_t row, float v0, float v1) {
+            if (row >= m) return;
+            if constexpr (OutFp8) {
+                if (col + 1 < n) {
+                    *reinterpret_cast<unsigned short*>(
+                        out_fp8 + row * n + col) =
+                        static_cast<unsigned short>(__nv_cvt_float2_to_fp8x2(
+                            make_float2(v0 * output_scale, v1 * output_scale),
+                            __NV_SATFINITE, __NV_E4M3));
+                } else {
+                    out_fp8[row * n + col] = __nv_fp8_e4m3(v0 * output_scale);
+                }
+            } else {
+                out_bf16[row * n + col] = __float2bfloat16(v0 * scale);
+                if (col + 1 < n)
+                    out_bf16[row * n + col + 1] = __float2bfloat16(v1 * scale);
+            }
+        };
+#pragma unroll
+        for (int mt = 0; mt < 4; ++mt) {
+            const int64_t row0 = row_base + mt * 16;
+            float* tile_acc = acc + (nt * 4 + mt) * 4;
+            if (col < n) {
+                store_out(row0, tile_acc[0], tile_acc[1]);
+                store_out(row0 + 8, tile_acc[2], tile_acc[3]);
+            }
+        }
+    }
+}
+
+template <bool AddBias = false, bool TrackAmax = true>
+void launch_fused_fp8_gemm_fast(
+    const torch::Tensor& a, const torch::Tensor& b, torch::Tensor& out,
+    const torch::Tensor& bias, const torch::Tensor& scale_a,
+    const torch::Tensor& scale_b, torch::Tensor* amax_a,
+    torch::Tensor* amax_b, int64_t m, int64_t n, int64_t k,
+    cudaStream_t stream) {
+    dim3 grid((n + kFastBlockN - 1) / kFastBlockN,
+              (m + kFastBlockM - 1) / kFastBlockM);
+    const auto* bias_ptr = AddBias
+        ? reinterpret_cast<const __nv_bfloat16*>(bias.data_ptr())
+        : nullptr;
+    auto kernel = fused_fp8_gemm_fast_kernel<AddBias, TrackAmax>;
+    static bool attribute_set = false;
+    if (!attribute_set) {
+        C10_CUDA_CHECK(cudaFuncSetAttribute(
+            kernel, cudaFuncAttributeMaxDynamicSharedMemorySize,
+            kFastSmemBytes));
+        attribute_set = true;
+    }
+    kernel<<<grid, kWarps * 32, kFastSmemBytes, stream>>>(
+        reinterpret_cast<const __nv_bfloat16*>(a.data_ptr()),
+        reinterpret_cast<const __nv_bfloat16*>(b.data_ptr()),
+        reinterpret_cast<__nv_bfloat16*>(out.data_ptr()), bias_ptr,
+        scale_a.data_ptr<float>(), scale_b.data_ptr<float>(),
+        amax_a ? amax_a->data_ptr<float>() : nullptr,
+        amax_b ? amax_b->data_ptr<float>() : nullptr, m, n, k);
+}
+
+void check_fp8_device(const torch::Tensor& tensor) {
+    static std::mutex mutex;
+    static std::unordered_map<int, bool> supported;
+    const int device = tensor.device().index();
+    {
+        std::lock_guard<std::mutex> lock(mutex);
+        auto cached = supported.find(device);
+        if (cached != supported.end()) {
+            TORCH_CHECK(cached->second,
+                        "fused FP8 MMA requires compute capability 8.9 or newer");
+            return;
+        }
+    }
+
+    const auto* properties = at::cuda::getDeviceProperties(device);
+    const bool is_supported = properties->major > 8 ||
+                              (properties->major == 8 && properties->minor >= 9);
+    {
+        std::lock_guard<std::mutex> lock(mutex);
+        supported.emplace(device, is_supported);
+    }
+    TORCH_CHECK(is_supported,
+                "fused FP8 MMA requires compute capability 8.9 or newer");
+}
+
+void check_scale(const torch::Tensor& scale, const torch::Tensor& input,
+                 const char* name) {
+    TORCH_CHECK(scale.is_cuda() && scale.device() == input.device() &&
+                    scale.scalar_type() == torch::kFloat32 && scale.numel() == 1,
+                name, " must be a CUDA float32 scalar on the input device");
+}
+
+}  // namespace
+
+torch::Tensor fp8_mm(torch::Tensor a, torch::Tensor b, torch::Tensor sx,
+                     torch::Tensor sw) {
+    TORCH_CHECK(a.is_cuda() && b.is_cuda(), "CUDA tensors required");
+    TORCH_CHECK(a.scalar_type() == torch::kBFloat16 &&
+                    b.scalar_type() == torch::kBFloat16,
+                "a and b must be bf16");
+    TORCH_CHECK(a.dim() == 2 && b.dim() == 2, "a and b must be 2D");
+    TORCH_CHECK(a.device() == b.device(), "a and b must be on the same device");
+    TORCH_CHECK(a.size(1) == b.size(1), "inner dim mismatch");
+    check_scale(sx, a, "sx");
+    check_scale(sw, a, "sw");
+    check_fp8_device(a);
+    const at::cuda::OptionalCUDAGuard guard(a.device());
+    auto stream = at::cuda::getCurrentCUDAStream();
+
+    auto a_c = a.contiguous();
+    auto b_c = b.contiguous();
+    auto out = torch::empty({a_c.size(0), b_c.size(0)}, a_c.options());
+    torch::Tensor no_bias;
+    launch_fused_fp8_gemm_fast<false, false>(
+        a_c, b_c, out, no_bias, sx, sw, nullptr, nullptr,
+        a_c.size(0), b_c.size(0), a_c.size(1), stream.stream());
+    C10_CUDA_CHECK(cudaGetLastError());
+    return out;
+}
+
+torch::Tensor fp8_linear_forward_scaled(
+    torch::Tensor x, torch::Tensor w, torch::Tensor bias, torch::Tensor sx,
+    torch::Tensor sw, torch::Tensor sx_inv, torch::Tensor sw_inv,
+    torch::Tensor amax_x, torch::Tensor amax_w) {
+    TORCH_CHECK(x.is_cuda() && w.is_cuda(), "CUDA tensors required");
+    TORCH_CHECK(x.scalar_type() == torch::kBFloat16 &&
+                    w.scalar_type() == torch::kBFloat16,
+                "x and w must be bf16");
+    TORCH_CHECK(x.device() == w.device(), "x and w must be on the same device");
+    check_scale(sx, x, "sx");
+    check_scale(sw, x, "sw");
+    check_fp8_device(x);
+    const at::cuda::OptionalCUDAGuard guard(x.device());
+    auto stream = at::cuda::getCurrentCUDAStream();
+
+    auto x_c = x.reshape({-1, w.size(1)}).contiguous();
+    auto w_c = w.contiguous();
+    int64_t m = x_c.size(0), k = x_c.size(1), n = w_c.size(0);
+    TORCH_CHECK(w_c.dim() == 2 && w_c.size(1) == k, "inner dim mismatch");
+    C10_CUDA_CHECK(cudaMemsetAsync(amax_x.data_ptr<float>(), 0, sizeof(float),
+                                   stream.stream()));
+    C10_CUDA_CHECK(cudaMemsetAsync(amax_w.data_ptr<float>(), 0, sizeof(float),
+                                   stream.stream()));
+    auto out = torch::empty({m, n}, x_c.options());
+    if (bias.defined() && bias.numel() > 0) {
+        TORCH_CHECK(bias.is_cuda() && bias.device() == x.device() &&
+                        bias.scalar_type() == torch::kBFloat16 &&
+                        bias.numel() == n,
+                    "bias must be CUDA bf16 with shape [N]");
+        launch_fused_fp8_gemm_fast<true, true>(
+            x_c, w_c, out, bias, sx, sw, &amax_x, &amax_w,
+            m, n, k, stream.stream());
+    } else {
+        launch_fused_fp8_gemm_fast<false, true>(
+            x_c, w_c, out, bias, sx, sw, &amax_x, &amax_w,
+            m, n, k, stream.stream());
+    }
+    C10_CUDA_CHECK(cudaGetLastError());
+
+    (void)sx_inv;
+    (void)sw_inv;
+    std::vector<int64_t> shape(x.sizes().begin(), x.sizes().end() - 1);
+    shape.push_back(n);
+    return out.reshape(shape);
+}
+
+std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> fp8_linear_backward_scaled(
+    torch::Tensor g, torch::Tensor x, torch::Tensor w,
+    std::vector<int64_t> masks, torch::Tensor sg, torch::Tensor sw,
+    torch::Tensor sx, torch::Tensor sg_inv, torch::Tensor sw_inv,
+    torch::Tensor sx_inv, torch::Tensor amax_g) {
+    TORCH_CHECK(g.is_cuda() && x.is_cuda() && w.is_cuda(), "CUDA tensors required");
+    TORCH_CHECK(g.scalar_type() == torch::kBFloat16 &&
+                    x.scalar_type() == torch::kBFloat16 &&
+                    w.scalar_type() == torch::kBFloat16,
+                "g, x, and w must be bf16");
+    TORCH_CHECK(g.device() == x.device() && g.device() == w.device(),
+                "g, x, and w must be on the same device");
+    TORCH_CHECK(masks.size() == 3, "masks must contain three values");
+    check_fp8_device(g);
+    const at::cuda::OptionalCUDAGuard guard(g.device());
+    auto stream = at::cuda::getCurrentCUDAStream();
+
+    auto g_c = g.reshape({-1, w.size(0)}).contiguous();
+    auto x_c = x.reshape({-1, x.size(-1)}).contiguous();
+    auto w_c = w.contiguous();
+    int64_t m = g_c.size(0), n = w_c.size(0), k = w_c.size(1);
+    TORCH_CHECK(x_c.size(0) == m && x_c.size(1) == k && g_c.size(1) == n,
+                "backward shape mismatch");
+
+    auto grad_input = torch::empty_like(x);
+    auto grad_weight = torch::empty_like(w);
+    auto grad_bias = torch::empty({0}, g.options());
+    C10_CUDA_CHECK(cudaMemsetAsync(amax_g.data_ptr<float>(), 0, sizeof(float),
+                                   stream.stream()));
+    torch::Tensor no_bias;
+    bool recorded_amax = false;
+    if (masks[0]) {
+        auto grad_input_2d = grad_input.reshape({m, k});
+        // The fast kernel computes A @ B^T. A contiguous W^T makes dX use
+        // the same coalesced forward tile path instead of scalar fragments.
+        auto w_t = w_c.transpose(0, 1).contiguous();
+        launch_fused_fp8_gemm_fast<false, true>(
+            g_c, w_t, grad_input_2d, no_bias, sg, sw, &amax_g, nullptr,
+            m, k, n, stream.stream());
+        recorded_amax = true;
+    }
+    if (masks[1]) {
+        // dW = G^T @ X, expressed as (G^T) @ (X^T)^T for the same kernel.
+        auto g_t = g_c.transpose(0, 1).contiguous();
+        auto x_t = x_c.transpose(0, 1).contiguous();
+        if (recorded_amax) {
+            launch_fused_fp8_gemm_fast<false, false>(
+                g_t, x_t, grad_weight, no_bias, sg, sx, nullptr, nullptr,
+                n, k, m, stream.stream());
+        } else {
+            launch_fused_fp8_gemm_fast<false, true>(
+                g_t, x_t, grad_weight, no_bias, sg, sx, &amax_g, nullptr,
+                n, k, m, stream.stream());
+        }
+        recorded_amax = true;
+    }
+    if (!recorded_amax) {
+        amax_g.copy_(g_c.abs().amax().to(torch::kFloat32));
+    }
+    C10_CUDA_CHECK(cudaGetLastError());
+    if (masks[2]) grad_bias = g_c.sum(0).to(g.scalar_type());
+
+    (void)sg_inv;
+    (void)sw_inv;
+    (void)sx_inv;
+    return {grad_input, grad_weight, grad_bias};
+}
+
+torch::Tensor fp8_mm_prequant(torch::Tensor a, torch::Tensor b,
+                              torch::Tensor scale) {
+    TORCH_CHECK(a.is_cuda() && b.is_cuda(), "CUDA tensors required");
+    TORCH_CHECK(a.scalar_type() == torch::kFloat8_e4m3fn &&
+                    b.scalar_type() == torch::kFloat8_e4m3fn,
+                "a and b must be fp8_e4m3fn");
+    TORCH_CHECK(a.dim() == 2 && b.dim() == 2, "a and b must be 2D");
+    TORCH_CHECK(a.device() == b.device(), "a and b must be on the same device");
+    TORCH_CHECK(a.size(1) == b.size(1), "inner dim mismatch");
+    check_scale(scale, a, "scale");
+    check_fp8_device(a);
+    const at::cuda::OptionalCUDAGuard guard(a.device());
+    auto stream = at::cuda::getCurrentCUDAStream();
+
+    auto a_c = a.contiguous();
+    auto b_c = b.contiguous();
+    int64_t m = a_c.size(0), k = a_c.size(1), n = b_c.size(0);
+    auto out = torch::empty({m, n},
+                            a_c.options().dtype(torch::kBFloat16));
+    const float scale_value = scale.item<float>();
+    dim3 grid((n + kPqBlockN - 1) / kPqBlockN,
+              (m + kPqBlockM - 1) / kPqBlockM);
+    fp8_mm_pq_kernel<false><<<grid, kWarps * 32, 0, stream>>>(
+        reinterpret_cast<const __nv_fp8_e4m3*>(a_c.data_ptr()),
+        reinterpret_cast<const __nv_fp8_e4m3*>(b_c.data_ptr()),
+        reinterpret_cast<__nv_bfloat16*>(out.data_ptr()), nullptr,
+        scale_value, 1.0f, m, n, k);
+    C10_CUDA_CHECK(cudaGetLastError());
+    return out;
+}
+
+torch::Tensor fp8_mm_prequant_fp8(torch::Tensor a, torch::Tensor b,
+                                  torch::Tensor scale,
+                                  torch::Tensor out_scale) {
+    TORCH_CHECK(a.is_cuda() && b.is_cuda(), "CUDA tensors required");
+    TORCH_CHECK(a.scalar_type() == torch::kFloat8_e4m3fn &&
+                    b.scalar_type() == torch::kFloat8_e4m3fn,
+                "a and b must be fp8_e4m3fn");
+    TORCH_CHECK(a.dim() == 2 && b.dim() == 2, "a and b must be 2D");
+    TORCH_CHECK(a.device() == b.device(), "a and b must be on the same device");
+    TORCH_CHECK(a.size(1) == b.size(1), "inner dim mismatch");
+    check_scale(scale, a, "scale");
+    check_scale(out_scale, a, "out_scale");
+    check_fp8_device(a);
+    const at::cuda::OptionalCUDAGuard guard(a.device());
+    auto stream = at::cuda::getCurrentCUDAStream();
+
+    auto a_c = a.contiguous();
+    auto b_c = b.contiguous();
+    int64_t m = a_c.size(0), k = a_c.size(1), n = b_c.size(0);
+    auto out = torch::empty({m, n}, a_c.options());
+    const float scale_value = scale.item<float>();
+    const float out_scale_value = out_scale.item<float>();
+    dim3 grid((n + kPqBlockN - 1) / kPqBlockN,
+              (m + kPqBlockM - 1) / kPqBlockM);
+    fp8_mm_pq_kernel<true><<<grid, kWarps * 32, 0, stream>>>(
+        reinterpret_cast<const __nv_fp8_e4m3*>(a_c.data_ptr()),
+        reinterpret_cast<const __nv_fp8_e4m3*>(b_c.data_ptr()), nullptr,
+        reinterpret_cast<__nv_fp8_e4m3*>(out.data_ptr()),
+        scale_value, out_scale_value, m, n, k);
+    C10_CUDA_CHECK(cudaGetLastError());
+    return out;
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("fp8_mm", &fp8_mm, py::arg("a"), py::arg("b"), py::arg("sx"),
+          py::arg("sw"),
+          "Fused BF16 input, E4M3 MMA, FP32 accumulation, BF16 output GEMM");
+    m.def("fp8_mm_prequant", &fp8_mm_prequant, py::arg("a"), py::arg("b"),
+          py::arg("scale"),
+          "Pre-quantized FP8 GEMM with FP32 accumulation and BF16 output");
+    m.def("fp8_mm_prequant_fp8", &fp8_mm_prequant_fp8, py::arg("a"),
+          py::arg("b"), py::arg("scale"), py::arg("out_scale"),
+          "Pre-quantized FP8 GEMM with FP32 accumulation and FP8 output");
+    m.def("fp8_linear_forward_scaled", &fp8_linear_forward_scaled,
+          py::arg("x"), py::arg("w"), py::arg("bias"), py::arg("sx"),
+          py::arg("sw"), py::arg("sx_inv"), py::arg("sw_inv"),
+          py::arg("amax_x"), py::arg("amax_w"),
+          "Fused BF16-to-FP8 linear forward with FP32 accumulation");
+    m.def("fp8_linear_backward_scaled", &fp8_linear_backward_scaled,
+          py::arg("g"), py::arg("x"), py::arg("w"), py::arg("masks"),
+          py::arg("sg"), py::arg("sw"), py::arg("sx"), py::arg("sg_inv"),
+          py::arg("sw_inv"), py::arg("sx_inv"), py::arg("amax_g"),
+          "Fused BF16-to-FP8 linear backward with FP32 accumulation");
+}

csrc/kernels/rotary_emb.cu

@@ -1,17 +1,18 @@
 #include <torch/extension.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <c10/cuda/CUDAException.h>
 #include <cuda_bf16.h>
 
 __global__ void rotary_emb_kernel(
     const __nv_bfloat16* __restrict__ x,
     const float* __restrict__ freqs_cis,
     __nv_bfloat16* __restrict__ out,
-    int batch,
-    int seq_len,
+    int n_tokens,
     int n_heads,
     int head_dim
 ) {
     const int half_dim = head_dim >> 1;
-    const int total = batch * seq_len * n_heads * half_dim;
+    const int total = n_tokens * n_heads * half_dim;
 
     for (int idx = blockIdx.x * blockDim.x + threadIdx.x;
          idx < total;
@@ -21,11 +22,10 @@ __global__ void rotary_emb_kernel(
         int tmp = idx / half_dim;
         int head = tmp % n_heads;
         tmp /= n_heads;
-        int seq = tmp % seq_len;
-        int b = tmp / seq_len;
+        int token = tmp;
 
-        int x_offset = ((b * seq_len + seq) * n_heads + head) * head_dim + (pair << 1);
-        int cs_offset = ((b * seq_len + seq) * half_dim + pair) * 2;
+        int x_offset = (token * n_heads + head) * head_dim + (pair << 1);
+        int cs_offset = (token * half_dim + pair) * 2;
 
         __nv_bfloat162 x_pair = *reinterpret_cast<const __nv_bfloat162*>(x + x_offset);
         float x_even = __bfloat162float(__low2bfloat16(x_pair));
@@ -46,34 +46,44 @@ torch::Tensor rotary_emb(
     torch::Tensor x,
     torch::Tensor freqs_cis
 ) {
+    const at::cuda::OptionalCUDAGuard device_guard(device_of(x));
+    auto stream = at::cuda::getCurrentCUDAStream();
+
     TORCH_CHECK(x.is_cuda(), "x must be on CUDA");
     TORCH_CHECK(freqs_cis.is_cuda(), "freqs_cis must be on CUDA");
     TORCH_CHECK(x.scalar_type() == torch::kBFloat16, "x must be bf16");
-    TORCH_CHECK(x.dim() == 4, "x must be 4D [batch, seq_len, n_heads, head_dim]");
+    TORCH_CHECK(x.dim() == 3 || x.dim() == 4,
+                "x must be [tokens, n_heads, head_dim] or "
+                "[batch, seq_len, n_heads, head_dim]");
     TORCH_CHECK(x.is_contiguous(), "x must be contiguous");
-    TORCH_CHECK(freqs_cis.dim() == 4, "freqs_cis must be 4D [batch, seq_len, dim/2, 2]");
+    TORCH_CHECK(freqs_cis.dim() == x.dim(), "freqs_cis rank must match x rank");
     TORCH_CHECK(freqs_cis.is_contiguous(), "freqs_cis must be contiguous");
+    TORCH_CHECK(freqs_cis.scalar_type() == torch::kFloat32, "freqs_cis must be f32");
 
-    int batch = x.size(0);
-    int seq_len = x.size(1);
-    int n_heads = x.size(2);
-    int head_dim = x.size(3);
+    int n_tokens = x.dim() == 3 ? x.size(0) : x.size(0) * x.size(1);
+    int n_heads = x.size(x.dim() - 2);
+    int head_dim = x.size(x.dim() - 1);
 
     TORCH_CHECK(head_dim % 2 == 0, "head_dim must be even");
+    TORCH_CHECK(freqs_cis.numel() == (int64_t)n_tokens * head_dim,
+                "freqs_cis token or rotary dimension mismatch");
+    TORCH_CHECK(freqs_cis.size(-2) == head_dim / 2, "freqs_cis dim/2 mismatch");
+    TORCH_CHECK(freqs_cis.size(-1) == 2, "freqs_cis last dim must be 2 [cos, sin]");
 
     auto out = torch::empty_like(x);
 
     int half_dim = head_dim / 2;
-    int total = batch * seq_len * n_heads * half_dim;
+    int total = n_tokens * n_heads * half_dim;
     int block = 256;
     int grid = std::min((total + block - 1) / block, 1024);
 
-    rotary_emb_kernel<<<grid, block>>>(
+    rotary_emb_kernel<<<grid, block, 0, stream>>>(
         reinterpret_cast<const __nv_bfloat16*>(x.data_ptr()),
         freqs_cis.data_ptr<float>(),
         reinterpret_cast<__nv_bfloat16*>(out.data_ptr()),
-        batch, seq_len, n_heads, head_dim
+        n_tokens, n_heads, head_dim
     );
+    C10_CUDA_CHECK(cudaGetLastError());
 
     return out;
 }
@@ -82,6 +92,6 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
     m.def("rotary_emb", &rotary_emb,
         py::arg("x"),
         py::arg("freqs_cis"),
-        "Fused rotary embedding (bf16 x, f32 freqs_cis [b,s,d/2,2], bf16 out)"
+        "Fused rotary embedding for packed 3D or dense 4D tensors"
     );
 }

csrc/tests/attn_decode_test.cu

@@ -1,185 +0,0 @@
-/*
-Pure-C test — uses shared dispatcher.
-nvcc -I csrc -arch=sm_89 -O3 \
-    --use_fast_math --ptxas-options=-O3 --extra-device-vectorization \
-    csrc/tests/attn_decode_test.cu -o test && ./test
-*/
-
-#include "test_utils.cuh"
-#include "../kernels/attn_dispatchers.cuh"
-
-// Split-K scratch (torch-free)
-struct DecodeScratch {
-    float* o_part = nullptr;
-    float* ml_part = nullptr;
-};
-
-static void setup_scratch(AttentionParams<bf16>& p, DecodeScratch& sc) {
-    int max_splits = 32;
-    cudaMalloc(&sc.o_part, (size_t)p.batch * p.q_head * max_splits * p.head_dim * sizeof(float));
-    cudaMalloc(&sc.ml_part, (size_t)p.batch * p.q_head * max_splits * 2 * sizeof(float));
-}
-
-static void free_scratch(DecodeScratch& sc) {
-    cudaFree(sc.o_part); cudaFree(sc.ml_part);
-}
-
-// Warmed-up, CUDA-event timed sweep over the production decode MMA path.
-static void bench() {
-    const int cfgs[][5] = {
-        {1, 32, 4, 512, 128},
-        {1, 32, 4, 1024, 128},
-        {1, 32, 4, 2048, 128},
-        {1, 32, 4, 4096, 128},
-        {16, 32, 4, 2048, 128},
-        {32, 32, 4, 1024, 128},
-    };
-    const int WARMUP = 10, ITERS = 100;
-    printf("\n===== DECODE BENCH (warmup=%d iters=%d) =====\n", WARMUP, ITERS);
-    print_bench_header();
-
-    for (int ci = 0; ci < 6; ci++) {
-        int B = cfgs[ci][0], Hq = cfgs[ci][1], Hk = cfgs[ci][2];
-        int sl = cfgs[ci][3], D = cfgs[ci][4];
-        size_t nQ = (size_t)B * Hq * D;
-        size_t nKV = (size_t)B * Hk * sl * D;
-
-        bf16 *dQ, *dK, *dV, *dO;
-        cudaMalloc(&dQ, nQ*2); cudaMalloc(&dK, nKV*2);
-        cudaMalloc(&dV, nKV*2); cudaMalloc(&dO, nQ*2);
-        size_t big = nQ > nKV ? nQ : nKV; bf16* tmp = new bf16[big];
-        for (size_t i = 0; i < nQ; i++)  tmp[i] = f2bf(randf());
-        cudaMemcpy(dQ, tmp, nQ*2, cudaMemcpyHostToDevice);
-        for (size_t i = 0; i < nKV; i++) tmp[i] = f2bf(randf());
-        cudaMemcpy(dK, tmp, nKV*2, cudaMemcpyHostToDevice);
-        for (size_t i = 0; i < nKV; i++) tmp[i] = f2bf(randf());
-        cudaMemcpy(dV, tmp, nKV*2, cudaMemcpyHostToDevice);
-        delete[] tmp;
-
-        AttentionParams<bf16> p;
-        p.batch = B; p.q_head = Hq; p.kv_head = Hk; p.q_len = 1; p.kv_len = sl;
-        p.head_dim = D; p.use_mask = 0; p.causal_offset = -1;
-        p.scale = 1.0f / sqrtf((float)D);
-        set_default_strides(p);
-        p.q = dQ; p.k = dK; p.v = dV; p.mask = nullptr; p.o = dO;
-
-        DecodeScratch sc;
-        setup_scratch(p, sc);
-        p.o_part = sc.o_part; p.ml_part = sc.ml_part;
-
-        auto launch = [&]() { dispatch_by_head_dim(D, [&]<int H>() { dispatch_decode<H>(p); }); };
-        double flops = 4.0 * B * Hq * (double)sl * D;
-        double bytes = 2.0 * (2.0 * nKV * sizeof(bf16));
-        BenchResult r = bench_kernel(launch, WARMUP, ITERS, flops, bytes);
-
-        char cfg[64];
-        snprintf(cfg, sizeof(cfg),
-                 "B=%2d Hq=%2d Hk=%d q=%4d kv=%4d D=%3d causal=%d",
-                 B, Hq, Hk, 1, sl, D, 0);
-        print_bench_row(cfg, r);
-
-        cudaFree(dQ); cudaFree(dK); cudaFree(dV); cudaFree(dO);
-        free_scratch(sc);
-    }
-}
-
-static int run_test(int B, int Hq, int Hk, int sl, int D, int causal) {
-    int gs = Hq / Hk;
-    printf("=== B=%d Hq=%d Hk=%d seq=%d D=%d gs=%d causal=%d ===\n",
-           B,Hq,Hk,sl,D,gs,causal);
-
-    size_t nQ = B*Hq*1*D, nKV = B*Hk*sl*D;
-    float *hQ=new float[nQ], *hK=new float[nKV], *hV=new float[nKV];
-    for (size_t i=0;i<nQ;i++) hQ[i]=randf();
-    for (size_t i=0;i<nKV;i++){hK[i]=randf();hV[i]=randf();}
-
-    bool* hMask=new bool[B*sl];
-    for (int i=0;i<B*sl;i++) hMask[i]=true;
-
-    bf16 *dQ,*dK,*dV,*dO,*tmp;
-    bool* dMask;
-    cudaMalloc(&dQ,nQ*2); cudaMalloc(&dK,nKV*2);
-    cudaMalloc(&dV,nKV*2); cudaMalloc(&dO,nQ*2);
-    cudaMalloc(&dMask,B*sl);
-
-    tmp=new bf16[max(nQ,nKV)];
-    for (size_t i=0;i<nQ;i++) tmp[i]=f2bf(hQ[i]);
-    cudaMemcpy(dQ,tmp,nQ*2,cudaMemcpyHostToDevice);
-    for (size_t i=0;i<nKV;i++) tmp[i]=f2bf(hK[i]);
-    cudaMemcpy(dK,tmp,nKV*2,cudaMemcpyHostToDevice);
-    for (size_t i=0;i<nKV;i++) tmp[i]=f2bf(hV[i]);
-    cudaMemcpy(dV,tmp,nKV*2,cudaMemcpyHostToDevice);
-    cudaMemcpy(dMask,hMask,B*sl,cudaMemcpyHostToDevice);
-
-    AttentionParams<bf16> p;
-    p.batch=B; p.q_head=Hq; p.kv_head=Hk; p.q_len=1; p.kv_len=sl; p.head_dim=D;
-    p.use_mask=0; p.causal_offset=causal?0:-1;
-    p.scale=1.0f/sqrtf((float)D);
-    set_default_strides(p);
-    p.q=dQ; p.k=dK; p.v=dV; p.mask=nullptr; p.o=dO;
-
-    DecodeScratch sc;
-    setup_scratch(p, sc);
-    p.o_part = sc.o_part; p.ml_part = sc.ml_part;
-
-    double t0=now_ms();
-    dispatch_by_head_dim(D, [&]<int H>() { dispatch_decode<H>(p); });
-    cudaDeviceSynchronize();
-    double kms=now_ms()-t0;
-    cudaError_t err=cudaGetLastError();
-    if (err!=cudaSuccess){printf("CUDA err: %s\n",cudaGetErrorString(err));return 1;}
-
-    bf16* hOut=new bf16[nQ];
-    cudaMemcpy(hOut,dO,nQ*2,cudaMemcpyDeviceToHost);
-
-    float* ref=new float[nQ];
-    cpu_attention_ref(hQ, hK, hV, hMask, ref, B, Hq, Hk, 1, sl, D, causal ? 0 : -1);
-
-    float max_abs_err=0, max_rel_err=0;
-    for (size_t i=0;i<nQ;i++){
-        float err=fabsf(bf2f(hOut[i])-ref[i]);
-        if(err>max_abs_err) max_abs_err=err;
-        float rel=err/fmaxf(fabsf(ref[i]), 1e-8f);
-        if(rel>max_rel_err) max_rel_err=rel;
-    }
-    const float atol=0.01f, rtol=0.01f;
-    bool pass=true;
-    for (size_t i=0;i<nQ;i++){
-        float err=fabsf(bf2f(hOut[i])-ref[i]);
-        if (err > atol + rtol * fabsf(ref[i])) { pass=false; break; }
-    }
-    printf("kernel: %.3f ms  max_abs_err: %.6e  max_rel_err: %.6e  %s\n\n",
-           kms, max_abs_err, max_rel_err, pass?"PASS":"FAIL");
-
-    cudaFree(dQ);cudaFree(dK);cudaFree(dV);cudaFree(dO);cudaFree(dMask);
-    free_scratch(sc);
-    delete[]hQ;delete[]hK;delete[]hV;delete[]hMask;delete[]hOut;delete[]ref;delete[]tmp;
-
-    return pass ? 0 : 1;
-}
-
-int main() {
-    const int configs[][6] = {
-        {1, 2, 1, 64, 32, 0},
-        {1, 32, 4, 512, 128, 0},
-        {1, 32, 4, 1024, 128, 0},
-        {1, 32, 4, 512, 128, 1},
-    };
-    int n_cfgs = sizeof(configs) / sizeof(configs[0]);
-    int fail = 0;
-
-    for (int ci = 0; ci < n_cfgs; ci++) {
-        int B = configs[ci][0], Hq = configs[ci][1], Hk = configs[ci][2];
-        int sl = configs[ci][3], D = configs[ci][4], causal = configs[ci][5];
-        fail += run_test(B, Hq, Hk, sl, D, causal);
-        if (fail) break;
-    }
-
-    if (fail) {
-        printf("FAILED\n");
-        return fail;
-    }
-    printf("All tests passed!\n");
-    bench();
-    return 0;
-}

csrc/tests/attn_paged_decode_test.cu

@@ -1,308 +0,0 @@
-// Compile:
-//   nvcc -I csrc -arch=sm_89 -O3 --use_fast_math --ptxas-options=-O3 \
-//        --extra-device-vectorization csrc/tests/attn_paged_decode_test.cu \
-//        -o /tmp/test_paged && /tmp/test_paged
-
-#include <cstring>
-#include "test_utils.cuh"
-#include "../kernels/attn_dispatchers.cuh"
-
-static void gather_kv_cpu(
-    const bf16* h_k_pool, const bf16* h_v_pool,
-    const int64_t* h_pt, int B, int Hkv, int kv_len,
-    int page_size, int head_dim,
-    bf16* h_k, bf16* h_v)
-{
-    int max_pages = (kv_len + page_size - 1) / page_size;
-    size_t page_stride = (size_t)page_size * Hkv * head_dim;
-    for (int b = 0; b < B; b++) {
-        for (int pos = 0; pos < kv_len; pos++) {
-            int log_pg = pos / page_size;
-            int pg_off = pos % page_size;
-            int phys = (int)h_pt[b * max_pages + log_pg];
-            for (int h = 0; h < Hkv; h++) {
-                size_t src_base = (size_t)phys * page_stride
-                                + (size_t)pg_off * Hkv * head_dim
-                                + h * head_dim;
-                size_t dst_base = ((size_t)b * Hkv + h) * kv_len * head_dim
-                                  + (size_t)pos * head_dim;
-                memcpy(h_k + dst_base, h_k_pool + src_base, head_dim * sizeof(bf16));
-                memcpy(h_v + dst_base, h_v_pool + src_base, head_dim * sizeof(bf16));
-            }
-        }
-    }
-}
-
-template <int HEAD_DIM>
-static int run_test(int B, int Hq, int Hkv, int kv_len, int page_size, int causal, int seed) {
-    printf("B=%d Hq=%d Hkv=%d kv_len=%d page_sz=%d head_dim=%d causal=%d ... ",
-           B, Hq, Hkv, kv_len, page_size, HEAD_DIM, causal);
-    fflush(stdout);
-
-    int max_pages = (kv_len + page_size - 1) / page_size;
-    int n_phys_pages = B * max_pages;
-    int max_splits = 32;
-
-    size_t sz_q  = (size_t)B * Hq * 1 * HEAD_DIM * sizeof(bf16);
-    size_t sz_o  = sz_q;
-    size_t sz_kv = (size_t)n_phys_pages * page_size * Hkv * HEAD_DIM * sizeof(bf16);
-    size_t sz_pt = (size_t)B * max_pages * sizeof(int64_t);
-    size_t sz_op = (size_t)B * Hq * max_splits * HEAD_DIM * sizeof(float);
-    size_t sz_ml = (size_t)B * Hq * max_splits * 2 * sizeof(float);
-
-    bf16 *d_q, *d_o_paged;
-    bf16 *d_k_pool, *d_v_pool;
-    int64_t* d_pt;
-    float *d_op, *d_ml;
-
-    cudaMalloc(&d_q, sz_q);
-    cudaMalloc(&d_o_paged, sz_o);
-    cudaMalloc(&d_k_pool, sz_kv);
-    cudaMalloc(&d_v_pool, sz_kv);
-    cudaMalloc(&d_pt, sz_pt);
-    cudaMalloc(&d_op, sz_op);
-    cudaMalloc(&d_ml, sz_ml);
-
-    srand(seed);
-    auto rnd = [&]() { return (rand() / (float)RAND_MAX) * 2.0f - 1.0f; };
-
-    bf16* h_q = (bf16*)malloc(sz_q);
-    for (int i = 0; i < B * Hq * HEAD_DIM; i++)
-        h_q[i] = __float2bfloat16(rnd());
-    cudaMemcpy(d_q, h_q, sz_q, cudaMemcpyHostToDevice);
-
-    bf16* h_k_pool = (bf16*)malloc(sz_kv);
-    bf16* h_v_pool = (bf16*)malloc(sz_kv);
-    size_t ps = (size_t)page_size * Hkv * HEAD_DIM;
-    for (int pg = 0; pg < n_phys_pages; pg++) {
-        for (int off = 0; off < page_size; off++) {
-            for (int h = 0; h < Hkv; h++) {
-                for (int d = 0; d < HEAD_DIM; d++) {
-                    float v = sinf((float)(pg * 7919 + off * 1049 + h * 331 + d));
-                    size_t idx = (size_t)pg * ps + (size_t)off * Hkv * HEAD_DIM
-                                 + h * HEAD_DIM + d;
-                    h_k_pool[idx] = __float2bfloat16(v);
-                    h_v_pool[idx] = __float2bfloat16(v * 0.3f);
-                }
-            }
-        }
-    }
-    cudaMemcpy(d_k_pool, h_k_pool, sz_kv, cudaMemcpyHostToDevice);
-    cudaMemcpy(d_v_pool, h_v_pool, sz_kv, cudaMemcpyHostToDevice);
-
-    int64_t* h_pt = (int64_t*)malloc(sz_pt);
-    int next_pg = 0;
-    for (int b = 0; b < B; b++)
-        for (int p = 0; p < max_pages; p++)
-            h_pt[b * max_pages + p] = next_pg++;
-    cudaMemcpy(d_pt, h_pt, sz_pt, cudaMemcpyHostToDevice);
-
-    bf16* h_k_cont = (bf16*)malloc((size_t)B * kv_len * Hkv * HEAD_DIM * sizeof(bf16));
-    bf16* h_v_cont = (bf16*)malloc((size_t)B * kv_len * Hkv * HEAD_DIM * sizeof(bf16));
-    gather_kv_cpu(h_k_pool, h_v_pool, h_pt, B, Hkv, kv_len, page_size, HEAD_DIM, h_k_cont, h_v_cont);
-
-    float* h_q_f = (float*)malloc((size_t)B * Hq * HEAD_DIM * sizeof(float));
-    float* h_k_f = (float*)malloc((size_t)B * kv_len * Hkv * HEAD_DIM * sizeof(float));
-    float* h_v_f = (float*)malloc((size_t)B * kv_len * Hkv * HEAD_DIM * sizeof(float));
-    for (int i = 0; i < B * Hq * HEAD_DIM; i++) h_q_f[i] = bf2f(h_q[i]);
-    for (int i = 0; i < B * kv_len * Hkv * HEAD_DIM; i++) {
-        h_k_f[i] = bf2f(h_k_cont[i]);
-        h_v_f[i] = bf2f(h_v_cont[i]);
-    }
-
-    float* h_o_ref = (float*)calloc(B * Hq * HEAD_DIM, sizeof(float));
-    cpu_attention_ref(h_q_f, h_k_f, h_v_f, nullptr, h_o_ref, B, Hq, Hkv,
-                      1, kv_len, HEAD_DIM, causal ? 0 : -1);
-
-    PagedAttentionParams<bf16> p;
-    p.batch = B; p.q_head = Hq; p.kv_head = Hkv; p.q_len = 1;
-    p.kv_len = kv_len; p.head_dim = HEAD_DIM;
-    p.use_mask = 0; p.causal_offset = causal ? 0 : -1;
-    set_default_paged_strides(p);
-    p.scale = 1.0f / sqrtf((float)HEAD_DIM);
-    p.page_size = page_size; p.max_pages = max_pages;
-    p.page_table = d_pt;
-    p.k_cache = d_k_pool; p.v_cache = d_v_pool;
-    p.q = d_q; p.mask = nullptr; p.o = d_o_paged;
-    p.o_part = d_op; p.ml_part = d_ml;
-
-    dispatch_by_head_dim(HEAD_DIM, [&]<int H>() { dispatch_paged_decode<H>(p); });
-    cudaDeviceSynchronize();
-
-    bf16* h_o_bf16 = (bf16*)malloc(sz_o);
-    cudaMemcpy(h_o_bf16, d_o_paged, sz_o, cudaMemcpyDeviceToHost);
-    float* h_o_paged = (float*)malloc(B * Hq * HEAD_DIM * sizeof(float));
-    for (int i = 0; i < B * Hq * HEAD_DIM; i++)
-        h_o_paged[i] = __bfloat162float(h_o_bf16[i]);
-
-    float max_abs_err = 0.0f, max_rel_err = 0.0f;
-    int bad_idx = -1;
-    for (int i = 0; i < B * Hq * HEAD_DIM; i++) {
-        float e = fabsf(h_o_paged[i] - h_o_ref[i]);
-        if (e > max_abs_err) { max_abs_err = e; bad_idx = i; }
-        float rel = e / fmaxf(fabsf(h_o_ref[i]), 1e-8f);
-        if (rel > max_rel_err) max_rel_err = rel;
-    }
-
-    const float atol = 0.01f, rtol = 0.01f;
-    bool pass = true;
-    for (int i = 0; i < B * Hq * HEAD_DIM; i++) {
-        float e = fabsf(h_o_paged[i] - h_o_ref[i]);
-        if (e > atol + rtol * fabsf(h_o_ref[i])) { pass = false; break; }
-    }
-
-    if (pass) {
-        printf("PASS (max_abs_err=%.4e max_rel_err=%.4e)\n", max_abs_err, max_rel_err);
-    } else {
-        int b = bad_idx / (Hq * HEAD_DIM);
-        int h = (bad_idx / HEAD_DIM) % Hq;
-        int d = bad_idx % HEAD_DIM;
-        printf("FAIL (max_abs_err=%.4e max_rel_err=%.4e at [%d,%d,%d]: ref=%.4f got=%.4f)\n",
-               max_abs_err, max_rel_err, b, h, d, h_o_ref[bad_idx], h_o_paged[bad_idx]);
-        printf("  ref[0..7]:");
-        for (int i = 0; i < 8 && i < HEAD_DIM; i++)
-            printf(" %.4f", h_o_ref[i]);
-        printf("\n  got[0..7]:");
-        for (int i = 0; i < 8 && i < HEAD_DIM; i++)
-            printf(" %.4f", h_o_paged[i]);
-        printf("\n");
-    }
-
-    free(h_q); free(h_k_pool); free(h_v_pool); free(h_pt);
-    free(h_k_cont); free(h_v_cont);
-    free(h_q_f); free(h_k_f); free(h_v_f);
-    free(h_o_ref); free(h_o_bf16); free(h_o_paged);
-    cudaFree(d_q); cudaFree(d_o_paged);
-    cudaFree(d_k_pool); cudaFree(d_v_pool); cudaFree(d_pt);
-    cudaFree(d_op); cudaFree(d_ml);
-
-    return pass ? 0 : 1;
-}
-
-struct TestCase {
-    int head_dim;
-    int B, Hq, Hkv, kv_len, page_size, causal, seed;
-};
-
-static const TestCase TESTS[] = {
-    {128, 1, 1, 1, 8, 128, 0, 1},
-    {128, 1, 4, 4, 128, 128, 0, 2},
-    {128, 2, 4, 4, 256, 128, 0, 3},
-    {128, 1, 4, 1, 64, 64, 0, 4},
-    {128, 1, 8, 2, 64, 128, 0, 5},
-    {128, 2, 16, 4, 128, 128, 0, 6},
-    {64, 1, 4, 2, 32, 128, 0, 7},
-    {256, 1, 2, 1, 16, 128, 0, 8},
-    {32, 1, 4, 2, 32, 64, 0, 9},
-    {128, 3, 8, 2, 256, 128, 0, 10},
-    {128, 2, 32, 8, 512, 128, 0, 11},
-    {128, 1, 16, 2, 256, 128, 0, 12},
-    {128, 2, 32, 4, 512, 128, 0, 13},
-    {128, 2, 8, 2, 128, 128, 1, 14},   // causal
-};
-
-static int dispatch_test(const TestCase& tc) {
-    int r = 0;
-    dispatch_by_head_dim(tc.head_dim, [&]<int D>() {
-        r = run_test<D>(tc.B, tc.Hq, tc.Hkv, tc.kv_len, tc.page_size, tc.causal, tc.seed);
-    });
-    return r;
-}
-
-template <int HEAD_DIM>
-static void bench_config(int B, int Hq, int Hkv, int kv_len, int page_size) {
-    int max_pages = (kv_len + page_size - 1) / page_size;
-    int n_phys_pages = B * max_pages;
-    int max_splits = 32;
-
-    size_t sz_q  = (size_t)B * Hq * 1 * HEAD_DIM * sizeof(bf16);
-    size_t sz_kv = (size_t)n_phys_pages * page_size * Hkv * HEAD_DIM * sizeof(bf16);
-    size_t sz_pt = (size_t)B * max_pages * sizeof(int64_t);
-    size_t sz_op = (size_t)B * Hq * max_splits * HEAD_DIM * sizeof(float);
-    size_t sz_ml = (size_t)B * Hq * max_splits * 2 * sizeof(float);
-
-    bf16 *d_q, *d_o, *d_k_pool, *d_v_pool;
-    int64_t* d_pt;
-    float *d_op, *d_ml;
-    cudaMalloc(&d_q, sz_q); cudaMalloc(&d_o, sz_q);
-    cudaMalloc(&d_k_pool, sz_kv); cudaMalloc(&d_v_pool, sz_kv);
-    cudaMalloc(&d_pt, sz_pt);
-    cudaMalloc(&d_op, sz_op); cudaMalloc(&d_ml, sz_ml);
-
-    bf16* tmp = (bf16*)malloc(sz_kv > sz_q ? sz_kv : sz_q);
-    for (size_t i = 0; i < sz_q / sizeof(bf16); i++) tmp[i] = f2bf(randf());
-    cudaMemcpy(d_q, tmp, sz_q, cudaMemcpyHostToDevice);
-    for (size_t i = 0; i < sz_kv / sizeof(bf16); i++) tmp[i] = f2bf(randf());
-    cudaMemcpy(d_k_pool, tmp, sz_kv, cudaMemcpyHostToDevice);
-    cudaMemcpy(d_v_pool, tmp, sz_kv, cudaMemcpyHostToDevice);
-
-    int64_t* h_pt = (int64_t*)malloc(sz_pt);
-    int next_pg = 0;
-    for (int b = 0; b < B; b++)
-        for (int p = 0; p < max_pages; p++)
-            h_pt[b * max_pages + p] = next_pg++;
-    cudaMemcpy(d_pt, h_pt, sz_pt, cudaMemcpyHostToDevice);
-    free(h_pt);
-
-    PagedAttentionParams<bf16> pa;
-    pa.batch = B; pa.q_head = Hq; pa.kv_head = Hkv; pa.q_len = 1;
-    pa.kv_len = kv_len; pa.head_dim = HEAD_DIM;
-    pa.use_mask = 0; pa.causal_offset = -1;
-    set_default_paged_strides(pa);
-    pa.scale = 1.0f / sqrtf((float)HEAD_DIM);
-    pa.page_size = page_size; pa.max_pages = max_pages;
-    pa.page_table = d_pt;
-    pa.k_cache = d_k_pool; pa.v_cache = d_v_pool;
-    pa.q = d_q; pa.mask = nullptr; pa.o = d_o;
-    pa.o_part = d_op; pa.ml_part = d_ml;
-
-    const int WARMUP = 10, ITERS = 100;
-    auto launch = [&]() {
-        dispatch_by_head_dim(HEAD_DIM, [&]<int H>() { dispatch_paged_decode<H>(pa); });
-    };
-    double flops = 4.0 * B * Hq * (double)kv_len * HEAD_DIM;
-    size_t nKV = (size_t)B * Hkv * kv_len * HEAD_DIM;
-    double bytes = 2.0 * (2.0 * nKV * sizeof(bf16));
-    BenchResult r = bench_kernel(launch, WARMUP, ITERS, flops, bytes);
-
-    char cfg[64];
-    snprintf(cfg, sizeof(cfg),
-             "B=%2d Hq=%2d Hk=%d q=%4d kv=%4d D=%3d page=%3d",
-             B, Hq, Hkv, 1, kv_len, HEAD_DIM, page_size);
-    print_bench_row(cfg, r);
-
-    free(tmp);
-    cudaFree(d_q); cudaFree(d_o);
-    cudaFree(d_k_pool); cudaFree(d_v_pool); cudaFree(d_pt);
-    cudaFree(d_op); cudaFree(d_ml);
-}
-
-static void bench() {
-    printf("\n===== PAGED DECODE BENCH =====\n");
-    print_bench_header();
-    bench_config<128>(1, 32, 4, 512, 128);
-    bench_config<128>(1, 32, 4, 1024, 128);
-    bench_config<128>(1, 32, 4, 2048, 128);
-    bench_config<128>(1, 32, 4, 4096, 128);
-    bench_config<128>(16, 32, 4, 2048, 128);
-    bench_config<128>(32, 32, 4, 1024, 128);
-}
-
-int main() {
-    int n = sizeof(TESTS) / sizeof(TESTS[0]);
-    int fail = 0;
-    printf("=== Paged Decode vs CPU reference (%d cases) ===\n\n", n);
-
-    for (int i = 0; i < n; i++) {
-        fail += dispatch_test(TESTS[i]);
-        if (fail) break;
-    }
-
-    if (fail) {
-        printf("\nFAILED (%d/%d tests failed)\n", fail, n);
-        return fail;
-    }
-    printf("\nAll %d tests passed!\n", n);
-    bench();
-    return 0;
-}

csrc/tests/attn_paged_test.cu

@@ -0,0 +1,990 @@
+// Compile:
+//   nvcc -I csrc -arch=sm_89 -O3 --use_fast_math --ptxas-options=-O3 \
+//        --extra-device-vectorization -Xcompiler -fopenmp \
+//        csrc/tests/attn_paged_test.cu \
+//        -o /tmp/test_paged && /tmp/test_paged
+
+#include <cstring>
+#include <vector>
+#include "test_utils.cuh"
+#include "../kernels/attn_dispatchers.cuh"
+
+struct PagedDecodeDispatch { AttentionParams<bf16>& p; template<int H> void operator()() { dispatch_paged_decode<H>(p, 0); } };
+struct PagedPrefillDispatch { AttentionParams<bf16>& p; template<int H> void operator()() { dispatch_paged_prefill<H>(p, 0); } };
+
+static int make_q_tile_mapping(const std::vector<int>& q_lens,
+                               int** d_batch, int** d_tile) {
+    constexpr int ROWS = 64;
+    std::vector<int> h_batch;
+    std::vector<int> h_tile;
+    for (int b = 0; b < (int)q_lens.size(); ++b) {
+        int n_tiles = (q_lens[b] + ROWS - 1) / ROWS;
+        for (int tile = 0; tile < n_tiles; ++tile) {
+            h_batch.push_back(b);
+            h_tile.push_back(tile);
+        }
+    }
+    size_t bytes = h_batch.size() * sizeof(int);
+    cudaMalloc(d_batch, bytes);
+    cudaMalloc(d_tile, bytes);
+    cudaMemcpy(*d_batch, h_batch.data(), bytes, cudaMemcpyHostToDevice);
+    cudaMemcpy(*d_tile, h_tile.data(), bytes, cudaMemcpyHostToDevice);
+    return (int)h_batch.size();
+}
+
+// ---- CPU reference: paged decode with variable seq_lens ----
+// Q: [B, Hq, D], K/V pool: [pool_size, Hkv, D]
+// req_to_token: [num_reqs, max_ctx_len], req_pool_indices: [B]
+// kv_indptr: [B+1].  mask: [B, max_seq_len] bool (True=keep) or NULL.
+static void cpu_paged_decode_ref(
+    const float* Q, const float* K_pool, const float* V_pool,
+    const int* req_to_token, const int* req_pool_indices,
+    const int* kv_indptr, const bool* mask, int mask_b_stride,
+    int B, int Hq, int Hkv, int D, int max_ctx_len,
+    float* O)
+{
+    float scale = 1.0f / sqrtf((float)D);
+    int n_rep = Hq / Hkv;
+    for (int b = 0; b < B; b++) {
+        int seq_len = kv_indptr[b + 1] - kv_indptr[b];
+        int req_idx = req_pool_indices[b];
+        #pragma omp parallel for schedule(dynamic)
+        for (int h = 0; h < Hq; h++) {
+            int kv_h = h / n_rep;
+            float mv = -INFINITY, sv = 0.0f;
+            float accum[256] = {0.0f};
+            for (int kj = 0; kj < seq_len; kj++) {
+                if (mask && !mask[b * mask_b_stride + kj]) continue;
+                int slot = req_to_token[req_idx * max_ctx_len + kj];
+                float dot = 0.0f;
+                for (int d = 0; d < D; d++)
+                    dot += Q[(b * Hq + h) * D + d] *
+                           K_pool[slot * Hkv * D + kv_h * D + d];
+                dot *= scale;
+                float nm = fmaxf(mv, dot);
+                float a = expf(mv - nm);
+                float be = expf(dot - nm);
+                sv = sv * a + be;
+                for (int d = 0; d < D; d++)
+                    accum[d] = accum[d] * a +
+                               V_pool[slot * Hkv * D + kv_h * D + d] * be;
+                mv = nm;
+            }
+            float inv = 1.0f / sv;
+            for (int d = 0; d < D; d++)
+                O[(b * Hq + h) * D + d] = accum[d] * inv;
+        }
+    }
+}
+
+// ---- CPU reference: paged prefill with ragged batch ----
+// Q: [total_q, Hq, D], K/V pool: [pool_size, Hkv, D]
+// req_to_token: [num_reqs, max_ctx_len], req_pool_indices: [B]
+// kv_indptr: [B+1], qo_indptr: [B+1].
+// mask: [B, max_q_len, max_seq_len] bool (True=keep, q-local + kv-local
+//   positions) or NULL.  Used only when causal==0 to apply an arbitrary
+//   attention mask on top of the (unused) causal logic.
+static void cpu_paged_prefill_ref(
+    const float* Q, const float* K_pool, const float* V_pool,
+    const int* req_to_token, const int* req_pool_indices,
+    const int* kv_indptr, const int* qo_indptr,
+    const bool* mask, int mask_l_stride, int mask_kv_stride,
+    int B, int Hq, int Hkv, int D, int max_ctx_len, int causal,
+    float* O)
+{
+    float scale = 1.0f / sqrtf((float)D);
+    int n_rep = Hq / Hkv;
+    for (int b = 0; b < B; b++) {
+        int seq_len = kv_indptr[b + 1] - kv_indptr[b];
+        int q_len = qo_indptr[b + 1] - qo_indptr[b];
+        int causal_off = seq_len - q_len;
+        int req_idx = req_pool_indices[b];
+        #pragma omp parallel for collapse(2) schedule(dynamic)
+        for (int h = 0; h < Hq; h++) {
+            for (int qi = 0; qi < q_len; qi++) {
+                int kv_h = h / n_rep;
+                float mv = -INFINITY, sv = 0.0f;
+                float accum[256] = {0.0f};
+                int lim = causal ? min(seq_len, causal_off + qi + 1) : seq_len;
+                for (int kj = 0; kj < lim; kj++) {
+                    if (mask && !mask[b * mask_l_stride * mask_kv_stride
+                                     + qi * mask_kv_stride + kj]) continue;
+                    int slot = req_to_token[req_idx * max_ctx_len + kj];
+                    float dot = 0.0f;
+                    for (int d = 0; d < D; d++)
+                        dot += Q[(qo_indptr[b] + qi) * Hq * D + h * D + d] *
+                               K_pool[slot * Hkv * D + kv_h * D + d];
+                    dot *= scale;
+                    float nm = fmaxf(mv, dot);
+                    float a = expf(mv - nm);
+                    float be = expf(dot - nm);
+                    sv = sv * a + be;
+                    for (int d = 0; d < D; d++)
+                        accum[d] = accum[d] * a +
+                                   V_pool[slot * Hkv * D + kv_h * D + d] * be;
+                    mv = nm;
+                }
+                float inv = 1.0f / sv;
+                for (int d = 0; d < D; d++)
+                    O[(qo_indptr[b] + qi) * Hq * D + h * D + d] = accum[d] * inv;
+            }
+        }
+    }
+}
+
+// ---- paged validation table (kernel vs CPU ref, abs error only) ----
+inline void print_paged_header() {
+    printf("%-58s | %11s | %6s\n",
+           "config", "max_err", "result");
+    printf("----------------------------------------------------------------"
+           "--------------------------------\n");
+}
+
+inline void print_paged_row(const char* cfg, float max_err, bool pass) {
+    printf("%-58s | %11.3e | %s\n",
+           cfg, max_err, pass ? "PASS" : "FAIL");
+}
+
+// ======================================================================
+// DECODE TEST
+// ======================================================================
+template <int HEAD_DIM>
+static int run_decode_test(int B, int Hq, int Hkv, int max_seq,
+                            int causal, int seed, int context_capacity = 0,
+                            int fixed_seq_len = 0) {
+    // Variable seq_lens per request
+    srand(seed);
+    std::vector<int> seq_lens(B);
+    for (int b = 0; b < B; b++)
+        seq_lens[b] = fixed_seq_len ? fixed_seq_len : 8 + rand() % (max_seq - 8);
+    int max_sl = *std::max_element(seq_lens.begin(), seq_lens.end());
+    int max_ctx = context_capacity ? context_capacity : max_sl + 16;
+
+    int pool_size = B * max_ctx;
+    int num_reqs = B + 4;
+
+    char cfg[80];
+    snprintf(cfg, sizeof(cfg), "DECODE B=%d Hq=%d Hkv=%d D=%d max_sl=%d causal=%d",
+             B, Hq, Hkv, HEAD_DIM, max_sl, causal);
+
+    size_t sz_q  = (size_t)B * Hq * HEAD_DIM * sizeof(bf16);
+    size_t sz_kv = (size_t)pool_size * Hkv * HEAD_DIM * sizeof(bf16);
+    size_t sz_rtt = (size_t)num_reqs * max_ctx * sizeof(int);
+    size_t sz_rpi = (size_t)B * sizeof(int);
+    size_t sz_kvi = (size_t)(B + 1) * sizeof(int);
+    size_t sz_op = (size_t)B * Hq * MAX_SPLITS * HEAD_DIM * sizeof(float);
+    size_t sz_ml = (size_t)B * Hq * MAX_SPLITS * 2 * sizeof(float);
+
+    bf16 *d_q, *d_o, *d_k_pool, *d_v_pool;
+    int *d_rtt, *d_rpi;
+    int *d_kvi;
+    float *d_op, *d_ml;
+    cudaMalloc(&d_q, sz_q); cudaMalloc(&d_o, sz_q);
+    cudaMalloc(&d_k_pool, sz_kv); cudaMalloc(&d_v_pool, sz_kv);
+    cudaMalloc(&d_rtt, sz_rtt); cudaMalloc(&d_rpi, sz_rpi);
+    cudaMalloc(&d_kvi, sz_kvi);
+    cudaMalloc(&d_op, sz_op); cudaMalloc(&d_ml, sz_ml);
+
+    auto rnd = [&]() { return (rand() / (float)RAND_MAX) * 2.0f - 1.0f; };
+
+    bf16* h_q = (bf16*)malloc(sz_q);
+    for (size_t i = 0; i < sz_q / sizeof(bf16); i++) h_q[i] = f2bf(rnd());
+    cudaMemcpy(d_q, h_q, sz_q, cudaMemcpyHostToDevice);
+
+    bf16* h_k_pool = (bf16*)malloc(sz_kv);
+    bf16* h_v_pool = (bf16*)malloc(sz_kv);
+    for (size_t i = 0; i < sz_kv / sizeof(bf16); i++) {
+        h_k_pool[i] = f2bf(rnd());
+        h_v_pool[i] = f2bf(rnd());
+    }
+    cudaMemcpy(d_k_pool, h_k_pool, sz_kv, cudaMemcpyHostToDevice);
+    cudaMemcpy(d_v_pool, h_v_pool, sz_kv, cudaMemcpyHostToDevice);
+
+    // req_to_token: assign unique slots per request (scattered, not contiguous)
+    int* h_rtt = (int*)malloc(sz_rtt);
+    int next_slot = 0;
+    for (int r = 0; r < num_reqs; r++)
+        for (int p = 0; p < max_ctx; p++) {
+            h_rtt[r * max_ctx + p] = next_slot % pool_size;
+            next_slot++;
+        }
+    cudaMemcpy(d_rtt, h_rtt, sz_rtt, cudaMemcpyHostToDevice);
+
+    // req_pool_indices: pick B random request rows
+    int* h_rpi = (int*)malloc(sz_rpi);
+    for (int b = 0; b < B; b++) h_rpi[b] = b;
+    cudaMemcpy(d_rpi, h_rpi, sz_rpi, cudaMemcpyHostToDevice);
+
+    // kv_indptr: prefix sum of seq_lens
+    int* h_kvi = (int*)malloc(sz_kvi);
+    h_kvi[0] = 0;
+    for (int b = 0; b < B; b++) h_kvi[b + 1] = h_kvi[b] + seq_lens[b];
+    cudaMemcpy(d_kvi, h_kvi, sz_kvi, cudaMemcpyHostToDevice);
+
+    // CPU reference
+    float* h_q_f = (float*)malloc(B * Hq * HEAD_DIM * sizeof(float));
+    float* h_k_f = (float*)malloc(pool_size * Hkv * HEAD_DIM * sizeof(float));
+    float* h_v_f = (float*)malloc(pool_size * Hkv * HEAD_DIM * sizeof(float));
+    for (int i = 0; i < B * Hq * HEAD_DIM; i++) h_q_f[i] = bf2f(h_q[i]);
+    for (int i = 0; i < pool_size * Hkv * HEAD_DIM; i++) {
+        h_k_f[i] = bf2f(h_k_pool[i]);
+        h_v_f[i] = bf2f(h_v_pool[i]);
+    }
+    float* h_o_ref = (float*)calloc(B * Hq * HEAD_DIM, sizeof(float));
+    cpu_paged_decode_ref(h_q_f, h_k_f, h_v_f, h_rtt, h_rpi, h_kvi,
+                            nullptr, 0,
+                            B, Hq, Hkv, HEAD_DIM, max_ctx, h_o_ref);
+
+    // Kernel launch
+    AttentionParams<bf16> p;
+    p.batch = B; p.q_head = Hq; p.kv_head = Hkv;
+    p.head_dim = HEAD_DIM;
+    p.q_l_stride = Hq * HEAD_DIM; p.q_h_stride = HEAD_DIM; p.q_d_stride = 1;
+    p.max_context_len = max_ctx;
+    p.causal_offset = causal ? 0 : -1; p.use_mask = 0;
+    p.mask = nullptr; p.mask_b_stride = 0;
+    p.mask_h_stride = 0; p.mask_l_stride = 0;
+    p.scale = 1.0f / sqrtf((float)HEAD_DIM);
+    p.q_ptr = d_q; p.k_ptr = d_k_pool; p.v_ptr = d_v_pool;
+    p.req_to_token = d_rtt; p.req_pool_indices = d_rpi;
+    p.kv_indptr = d_kvi; p.qo_indptr = nullptr;
+    p.o_ptr = d_o; p.o_part = d_op; p.ml_part = d_ml;
+
+    dispatch_by_head_dim(HEAD_DIM, PagedDecodeDispatch{p});
+    cudaDeviceSynchronize();
+
+    bf16* h_o_bf = (bf16*)malloc(sz_q);
+    cudaMemcpy(h_o_bf, d_o, sz_q, cudaMemcpyDeviceToHost);
+    float* h_o_got = (float*)malloc(B * Hq * HEAD_DIM * sizeof(float));
+    for (int i = 0; i < B * Hq * HEAD_DIM; i++) h_o_got[i] = bf2f(h_o_bf[i]);
+
+    const float atol = 0.01f, rtol = 0.01f;
+    bool pass = true;
+    float max_err = 0.0f;
+    for (int i = 0; i < B * Hq * HEAD_DIM; i++) {
+        float e = fabsf(h_o_got[i] - h_o_ref[i]);
+        if (e > max_err) max_err = e;
+        if (e > atol + rtol * fabsf(h_o_ref[i])) { pass = false; break; }
+    }
+
+    print_paged_row(cfg, max_err, pass);
+
+    free(h_q); free(h_k_pool); free(h_v_pool); free(h_rtt); free(h_rpi);
+    free(h_kvi); free(h_q_f); free(h_k_f); free(h_v_f);
+    free(h_o_ref); free(h_o_bf); free(h_o_got);
+    cudaFree(d_q); cudaFree(d_o); cudaFree(d_k_pool); cudaFree(d_v_pool);
+    cudaFree(d_rtt); cudaFree(d_rpi); cudaFree(d_kvi); cudaFree(d_op); cudaFree(d_ml);
+    return pass ? 0 : 1;
+}
+
+// ======================================================================
+// DECODE WITH MASK TEST (regression: 2D mask on mixed seq_lens)
+// ======================================================================
+template <int HEAD_DIM>
+static int run_decode_mask_test(int B, int Hq, int Hkv, int max_seq,
+                                int seed) {
+    srand(seed);
+    std::vector<int> seq_lens(B);
+    for (int b = 0; b < B; b++)
+        seq_lens[b] = 8 + rand() % (max_seq - 8);
+    int max_sl = *std::max_element(seq_lens.begin(), seq_lens.end());
+    int max_ctx = max_sl + 16;
+    int pool_size = B * max_ctx;
+    int num_reqs = B + 4;
+
+    char cfg[80];
+    snprintf(cfg, sizeof(cfg), "DECODE-MASK B=%d Hq=%d Hkv=%d D=%d max_sl=%d",
+             B, Hq, Hkv, HEAD_DIM, max_sl);
+
+    size_t sz_q  = (size_t)B * Hq * HEAD_DIM * sizeof(bf16);
+    size_t sz_kv = (size_t)pool_size * Hkv * HEAD_DIM * sizeof(bf16);
+    size_t sz_rtt = (size_t)num_reqs * max_ctx * sizeof(int);
+    size_t sz_rpi = (size_t)B * sizeof(int);
+    size_t sz_kvi = (size_t)(B + 1) * sizeof(int);
+    size_t sz_mask = (size_t)B * max_sl * sizeof(bool);
+    size_t sz_op = (size_t)B * Hq * MAX_SPLITS * HEAD_DIM * sizeof(float);
+    size_t sz_ml = (size_t)B * Hq * MAX_SPLITS * 2 * sizeof(float);
+
+    bf16 *d_q, *d_o, *d_k_pool, *d_v_pool;
+    int *d_rtt, *d_rpi;
+    int *d_kvi;
+    bool *d_mask;
+    float *d_op, *d_ml;
+    cudaMalloc(&d_q, sz_q); cudaMalloc(&d_o, sz_q);
+    cudaMalloc(&d_k_pool, sz_kv); cudaMalloc(&d_v_pool, sz_kv);
+    cudaMalloc(&d_rtt, sz_rtt); cudaMalloc(&d_rpi, sz_rpi);
+    cudaMalloc(&d_kvi, sz_kvi);
+    cudaMalloc(&d_mask, sz_mask);
+    cudaMalloc(&d_op, sz_op); cudaMalloc(&d_ml, sz_ml);
+
+    auto rnd = [&]() { return (rand() / (float)RAND_MAX) * 2.0f - 1.0f; };
+
+    bf16* h_q = (bf16*)malloc(sz_q);
+    for (size_t i = 0; i < sz_q / sizeof(bf16); i++) h_q[i] = f2bf(rnd());
+    cudaMemcpy(d_q, h_q, sz_q, cudaMemcpyHostToDevice);
+
+    bf16* h_k_pool = (bf16*)malloc(sz_kv);
+    bf16* h_v_pool = (bf16*)malloc(sz_kv);
+    for (size_t i = 0; i < sz_kv / sizeof(bf16); i++) {
+        h_k_pool[i] = f2bf(rnd());
+        h_v_pool[i] = f2bf(rnd());
+    }
+    cudaMemcpy(d_k_pool, h_k_pool, sz_kv, cudaMemcpyHostToDevice);
+    cudaMemcpy(d_v_pool, h_v_pool, sz_kv, cudaMemcpyHostToDevice);
+
+    int* h_rtt = (int*)malloc(sz_rtt);
+    int next_slot = 0;
+    for (int r = 0; r < num_reqs; r++)
+        for (int p = 0; p < max_ctx; p++) {
+            h_rtt[r * max_ctx + p] = next_slot % pool_size;
+            next_slot++;
+        }
+    cudaMemcpy(d_rtt, h_rtt, sz_rtt, cudaMemcpyHostToDevice);
+
+    int* h_rpi = (int*)malloc(sz_rpi);
+    for (int b = 0; b < B; b++) h_rpi[b] = b;
+    cudaMemcpy(d_rpi, h_rpi, sz_rpi, cudaMemcpyHostToDevice);
+
+    int* h_kvi = (int*)malloc(sz_kvi);
+    h_kvi[0] = 0;
+    for (int b = 0; b < B; b++) h_kvi[b + 1] = h_kvi[b] + seq_lens[b];
+    cudaMemcpy(d_kvi, h_kvi, sz_kvi, cudaMemcpyHostToDevice);
+
+    // Mask: keep first half of each request's kv range, drop the rest —
+    // exercises the HasMask path with per-request seq_len.
+    bool* h_mask = (bool*)malloc(sz_mask);
+    for (int b = 0; b < B; b++)
+        for (int k = 0; k < max_sl; k++)
+            h_mask[b * max_sl + k] = (k < seq_lens[b]) && (k % 2 == 0);
+    cudaMemcpy(d_mask, h_mask, sz_mask, cudaMemcpyHostToDevice);
+
+    float* h_q_f = (float*)malloc(B * Hq * HEAD_DIM * sizeof(float));
+    float* h_k_f = (float*)malloc(pool_size * Hkv * HEAD_DIM * sizeof(float));
+    float* h_v_f = (float*)malloc(pool_size * Hkv * HEAD_DIM * sizeof(float));
+    for (int i = 0; i < B * Hq * HEAD_DIM; i++) h_q_f[i] = bf2f(h_q[i]);
+    for (int i = 0; i < pool_size * Hkv * HEAD_DIM; i++) {
+        h_k_f[i] = bf2f(h_k_pool[i]);
+        h_v_f[i] = bf2f(h_v_pool[i]);
+    }
+    float* h_o_ref = (float*)calloc(B * Hq * HEAD_DIM, sizeof(float));
+    cpu_paged_decode_ref(h_q_f, h_k_f, h_v_f, h_rtt, h_rpi, h_kvi,
+                            h_mask, max_sl,
+                            B, Hq, Hkv, HEAD_DIM, max_ctx, h_o_ref);
+
+    AttentionParams<bf16> p;
+    p.batch = B; p.q_head = Hq; p.kv_head = Hkv;
+    p.head_dim = HEAD_DIM;
+    p.q_l_stride = Hq * HEAD_DIM; p.q_h_stride = HEAD_DIM; p.q_d_stride = 1;
+    p.max_context_len = max_ctx;
+    p.causal_offset = -1; p.use_mask = 1;
+    p.mask = d_mask; p.mask_b_stride = max_sl;
+    p.mask_h_stride = 0; p.mask_l_stride = 0;
+    p.scale = 1.0f / sqrtf((float)HEAD_DIM);
+    p.q_ptr = d_q; p.k_ptr = d_k_pool; p.v_ptr = d_v_pool;
+    p.req_to_token = d_rtt; p.req_pool_indices = d_rpi;
+    p.kv_indptr = d_kvi; p.qo_indptr = nullptr;
+    p.o_ptr = d_o; p.o_part = d_op; p.ml_part = d_ml;
+
+    dispatch_by_head_dim(HEAD_DIM, PagedDecodeDispatch{p});
+    cudaDeviceSynchronize();
+
+    bf16* h_o_bf = (bf16*)malloc(sz_q);
+    cudaMemcpy(h_o_bf, d_o, sz_q, cudaMemcpyDeviceToHost);
+    float* h_o_got = (float*)malloc(B * Hq * HEAD_DIM * sizeof(float));
+    for (int i = 0; i < B * Hq * HEAD_DIM; i++) h_o_got[i] = bf2f(h_o_bf[i]);
+
+    const float atol = 0.01f, rtol = 0.01f;
+    bool pass = true;
+    float max_err = 0.0f;
+    for (int i = 0; i < B * Hq * HEAD_DIM; i++) {
+        float e = fabsf(h_o_got[i] - h_o_ref[i]);
+        if (e > max_err) max_err = e;
+        if (e > atol + rtol * fabsf(h_o_ref[i])) { pass = false; break; }
+    }
+
+    print_paged_row(cfg, max_err, pass);
+
+    free(h_q); free(h_k_pool); free(h_v_pool); free(h_rtt); free(h_rpi);
+    free(h_kvi); free(h_mask); free(h_q_f); free(h_k_f); free(h_v_f);
+    free(h_o_ref); free(h_o_bf); free(h_o_got);
+    cudaFree(d_q); cudaFree(d_o); cudaFree(d_k_pool); cudaFree(d_v_pool);
+    cudaFree(d_rtt); cudaFree(d_rpi); cudaFree(d_kvi); cudaFree(d_mask);
+    cudaFree(d_op); cudaFree(d_ml);
+    return pass ? 0 : 1;
+}
+
+// ======================================================================
+// PREFILL TEST
+// ======================================================================
+template <int HEAD_DIM>
+static int run_prefill_test(int B, int Hq, int Hkv,
+                             std::vector<int>& q_lens,
+                             std::vector<int>& kv_lens,
+                             int causal, int seed) {
+    int total_q = 0;
+    int max_sl = 0;
+    for (int b = 0; b < B; b++) {
+        total_q += q_lens[b];
+        max_sl = max(max_sl, kv_lens[b]);
+    }
+    int max_ctx = max_sl + 16;
+    int pool_size = B * max_ctx;
+    int num_reqs = B + 4;
+
+    char cfg[80];
+    snprintf(cfg, sizeof(cfg), "PREFILL B=%d Hq=%d Hkv=%d D=%d max_sl=%d causal=%d",
+             B, Hq, Hkv, HEAD_DIM, max_sl, causal);
+
+    size_t sz_q  = (size_t)total_q * Hq * HEAD_DIM * sizeof(bf16);
+    size_t sz_kv = (size_t)pool_size * Hkv * HEAD_DIM * sizeof(bf16);
+    size_t sz_rtt = (size_t)num_reqs * max_ctx * sizeof(int);
+    size_t sz_rpi = (size_t)B * sizeof(int);
+    size_t sz_kvi = (size_t)(B + 1) * sizeof(int);
+    size_t sz_qoi = (size_t)(B + 1) * sizeof(int);
+
+    bf16 *d_q, *d_o, *d_k_pool, *d_v_pool;
+    int *d_rtt, *d_rpi;
+    int *d_kvi, *d_qoi;
+    cudaMalloc(&d_q, sz_q); cudaMalloc(&d_o, sz_q);
+    cudaMalloc(&d_k_pool, sz_kv); cudaMalloc(&d_v_pool, sz_kv);
+    cudaMalloc(&d_rtt, sz_rtt); cudaMalloc(&d_rpi, sz_rpi);
+    cudaMalloc(&d_kvi, sz_kvi); cudaMalloc(&d_qoi, sz_qoi);
+
+    srand(seed);
+    auto rnd = [&]() { return (rand() / (float)RAND_MAX) * 2.0f - 1.0f; };
+
+    bf16* h_q = (bf16*)malloc(sz_q);
+    for (size_t i = 0; i < sz_q / sizeof(bf16); i++) h_q[i] = f2bf(rnd());
+    cudaMemcpy(d_q, h_q, sz_q, cudaMemcpyHostToDevice);
+
+    bf16* h_k_pool = (bf16*)malloc(sz_kv);
+    bf16* h_v_pool = (bf16*)malloc(sz_kv);
+    for (size_t i = 0; i < sz_kv / sizeof(bf16); i++) {
+        h_k_pool[i] = f2bf(rnd());
+        h_v_pool[i] = f2bf(rnd());
+    }
+    cudaMemcpy(d_k_pool, h_k_pool, sz_kv, cudaMemcpyHostToDevice);
+    cudaMemcpy(d_v_pool, h_v_pool, sz_kv, cudaMemcpyHostToDevice);
+
+    int* h_rtt = (int*)malloc(sz_rtt);
+    int next_slot = 0;
+    for (int r = 0; r < num_reqs; r++)
+        for (int p = 0; p < max_ctx; p++) {
+            h_rtt[r * max_ctx + p] = next_slot % pool_size;
+            next_slot++;
+        }
+    cudaMemcpy(d_rtt, h_rtt, sz_rtt, cudaMemcpyHostToDevice);
+
+    int* h_rpi = (int*)malloc(sz_rpi);
+    for (int b = 0; b < B; b++) h_rpi[b] = b;
+    cudaMemcpy(d_rpi, h_rpi, sz_rpi, cudaMemcpyHostToDevice);
+
+    int* h_kvi = (int*)malloc(sz_kvi);
+    h_kvi[0] = 0;
+    for (int b = 0; b < B; b++) h_kvi[b + 1] = h_kvi[b] + kv_lens[b];
+    cudaMemcpy(d_kvi, h_kvi, sz_kvi, cudaMemcpyHostToDevice);
+
+    int* h_qoi = (int*)malloc(sz_qoi);
+    h_qoi[0] = 0;
+    for (int b = 0; b < B; b++) h_qoi[b + 1] = h_qoi[b] + q_lens[b];
+    cudaMemcpy(d_qoi, h_qoi, sz_qoi, cudaMemcpyHostToDevice);
+
+    // CPU reference
+    float* h_q_f = (float*)malloc(total_q * Hq * HEAD_DIM * sizeof(float));
+    float* h_k_f = (float*)malloc(pool_size * Hkv * HEAD_DIM * sizeof(float));
+    float* h_v_f = (float*)malloc(pool_size * Hkv * HEAD_DIM * sizeof(float));
+    for (int i = 0; i < total_q * Hq * HEAD_DIM; i++) h_q_f[i] = bf2f(h_q[i]);
+    for (int i = 0; i < pool_size * Hkv * HEAD_DIM; i++) {
+        h_k_f[i] = bf2f(h_k_pool[i]);
+        h_v_f[i] = bf2f(h_v_pool[i]);
+    }
+    float* h_o_ref = (float*)calloc(total_q * Hq * HEAD_DIM, sizeof(float));
+    cpu_paged_prefill_ref(h_q_f, h_k_f, h_v_f, h_rtt, h_rpi, h_kvi, h_qoi,
+                             nullptr, 0, 0,
+                             B, Hq, Hkv, HEAD_DIM, max_ctx, causal, h_o_ref);
+
+    int *d_qtb, *d_qti;
+    int num_q_tiles = make_q_tile_mapping(q_lens, &d_qtb, &d_qti);
+
+    // Kernel launch
+    AttentionParams<bf16> p;
+    p.batch = B; p.q_head = Hq; p.kv_head = Hkv;
+    p.head_dim = HEAD_DIM;
+    p.q_l_stride = Hq * HEAD_DIM; p.q_h_stride = HEAD_DIM; p.q_d_stride = 1;
+    p.max_context_len = max_ctx;
+    p.q_len = total_q;
+    p.causal_offset = causal ? 0 : -1; p.use_mask = 0;
+    p.mask = nullptr; p.mask_b_stride = 0;
+    p.mask_h_stride = 0; p.mask_l_stride = 0;
+    p.scale = 1.0f / sqrtf((float)HEAD_DIM);
+    p.q_ptr = d_q; p.k_ptr = d_k_pool; p.v_ptr = d_v_pool;
+    p.req_to_token = d_rtt; p.req_pool_indices = d_rpi;
+    p.kv_indptr = d_kvi; p.qo_indptr = d_qoi;
+    p.q_tile_to_batch = d_qtb; p.q_tile_to_index = d_qti;
+    p.num_q_tiles = num_q_tiles;
+    p.o_ptr = d_o; p.o_part = nullptr; p.ml_part = nullptr;
+
+    dispatch_by_head_dim(HEAD_DIM, PagedPrefillDispatch{p});
+    cudaDeviceSynchronize();
+
+    bf16* h_o_bf = (bf16*)malloc(sz_q);
+    cudaMemcpy(h_o_bf, d_o, sz_q, cudaMemcpyDeviceToHost);
+    float* h_o_got = (float*)malloc(total_q * Hq * HEAD_DIM * sizeof(float));
+    for (int i = 0; i < total_q * Hq * HEAD_DIM; i++) h_o_got[i] = bf2f(h_o_bf[i]);
+
+    const float atol = 0.01f, rtol = 0.01f;
+    bool pass = true;
+    float max_err = 0.0f;
+    for (int i = 0; i < total_q * Hq * HEAD_DIM; i++) {
+        float e = fabsf(h_o_got[i] - h_o_ref[i]);
+        if (e > max_err) max_err = e;
+        if (e > atol + rtol * fabsf(h_o_ref[i])) { pass = false; break; }
+    }
+
+    print_paged_row(cfg, max_err, pass);
+
+    free(h_q); free(h_k_pool); free(h_v_pool); free(h_rtt); free(h_rpi);
+    free(h_kvi); free(h_qoi); free(h_q_f); free(h_k_f); free(h_v_f);
+    free(h_o_ref); free(h_o_bf); free(h_o_got);
+    cudaFree(d_q); cudaFree(d_o); cudaFree(d_k_pool); cudaFree(d_v_pool);
+    cudaFree(d_rtt); cudaFree(d_rpi); cudaFree(d_kvi); cudaFree(d_qoi);
+    cudaFree(d_qtb); cudaFree(d_qti);
+    return pass ? 0 : 1;
+}
+
+// ======================================================================
+// PREFILL WITH MASK TEST (regression: 4D causal mask on single request)
+// ======================================================================
+template <int HEAD_DIM>
+static int run_prefill_mask_test(int Hq, int Hkv, int q_len, int seed) {
+    srand(seed);
+    int B = 1;
+    int total_q = q_len;
+    int seq_len = q_len;  // pure prefill: kv_len == q_len
+    int max_ctx = seq_len + 16;
+    int pool_size = B * max_ctx;
+    int num_reqs = B + 4;
+
+    char cfg[80];
+    snprintf(cfg, sizeof(cfg), "PREFILL-MASK Hq=%d Hkv=%d D=%d q_len=%d",
+             Hq, Hkv, HEAD_DIM, q_len);
+    fflush(stdout);
+
+    size_t sz_q  = (size_t)total_q * Hq * HEAD_DIM * sizeof(bf16);
+    size_t sz_kv = (size_t)pool_size * Hkv * HEAD_DIM * sizeof(bf16);
+    size_t sz_rtt = (size_t)num_reqs * max_ctx * sizeof(int);
+    size_t sz_rpi = (size_t)B * sizeof(int);
+    size_t sz_kvi = (size_t)(B + 1) * sizeof(int);
+    size_t sz_qoi = (size_t)(B + 1) * sizeof(int);
+    size_t sz_mask = (size_t)B * q_len * q_len * sizeof(bool);
+
+    bf16 *d_q, *d_o, *d_k_pool, *d_v_pool;
+    int *d_rtt, *d_rpi;
+    int *d_kvi, *d_qoi;
+    bool *d_mask;
+    cudaMalloc(&d_q, sz_q); cudaMalloc(&d_o, sz_q);
+    cudaMalloc(&d_k_pool, sz_kv); cudaMalloc(&d_v_pool, sz_kv);
+    cudaMalloc(&d_rtt, sz_rtt); cudaMalloc(&d_rpi, sz_rpi);
+    cudaMalloc(&d_kvi, sz_kvi); cudaMalloc(&d_qoi, sz_qoi);
+    cudaMalloc(&d_mask, sz_mask);
+
+    auto rnd = [&]() { return (rand() / (float)RAND_MAX) * 2.0f - 1.0f; };
+
+    bf16* h_q = (bf16*)malloc(sz_q);
+    for (size_t i = 0; i < sz_q / sizeof(bf16); i++) h_q[i] = f2bf(rnd());
+    cudaMemcpy(d_q, h_q, sz_q, cudaMemcpyHostToDevice);
+
+    bf16* h_k_pool = (bf16*)malloc(sz_kv);
+    bf16* h_v_pool = (bf16*)malloc(sz_kv);
+    for (size_t i = 0; i < sz_kv / sizeof(bf16); i++) {
+        h_k_pool[i] = f2bf(rnd());
+        h_v_pool[i] = f2bf(rnd());
+    }
+    cudaMemcpy(d_k_pool, h_k_pool, sz_kv, cudaMemcpyHostToDevice);
+    cudaMemcpy(d_v_pool, h_v_pool, sz_kv, cudaMemcpyHostToDevice);
+
+    int* h_rtt = (int*)malloc(sz_rtt);
+    int next_slot = 0;
+    for (int r = 0; r < num_reqs; r++)
+        for (int p = 0; p < max_ctx; p++) {
+            h_rtt[r * max_ctx + p] = next_slot % pool_size;
+            next_slot++;
+        }
+    cudaMemcpy(d_rtt, h_rtt, sz_rtt, cudaMemcpyHostToDevice);
+
+    int* h_rpi = (int*)malloc(sz_rpi);
+    h_rpi[0] = 0;
+    cudaMemcpy(d_rpi, h_rpi, sz_rpi, cudaMemcpyHostToDevice);
+
+    int* h_kvi = (int*)malloc(sz_kvi);
+    h_kvi[0] = 0; h_kvi[1] = seq_len;
+    cudaMemcpy(d_kvi, h_kvi, sz_kvi, cudaMemcpyHostToDevice);
+
+    int* h_qoi = (int*)malloc(sz_qoi);
+    h_qoi[0] = 0; h_qoi[1] = q_len;
+    cudaMemcpy(d_qoi, h_qoi, sz_qoi, cudaMemcpyHostToDevice);
+
+    // 4D causal mask [B, 1, q_len, q_len], True=keep.
+    bool* h_mask = (bool*)malloc(sz_mask);
+    for (int qi = 0; qi < q_len; qi++)
+        for (int kj = 0; kj < q_len; kj++)
+            h_mask[qi * q_len + kj] = (kj <= qi);
+    cudaMemcpy(d_mask, h_mask, sz_mask, cudaMemcpyHostToDevice);
+
+    float* h_q_f = (float*)malloc(total_q * Hq * HEAD_DIM * sizeof(float));
+    float* h_k_f = (float*)malloc(pool_size * Hkv * HEAD_DIM * sizeof(float));
+    float* h_v_f = (float*)malloc(pool_size * Hkv * HEAD_DIM * sizeof(float));
+    for (int i = 0; i < total_q * Hq * HEAD_DIM; i++) h_q_f[i] = bf2f(h_q[i]);
+    for (int i = 0; i < pool_size * Hkv * HEAD_DIM; i++) {
+        h_k_f[i] = bf2f(h_k_pool[i]);
+        h_v_f[i] = bf2f(h_v_pool[i]);
+    }
+    float* h_o_ref = (float*)calloc(total_q * Hq * HEAD_DIM, sizeof(float));
+    // CPU ref with causal=0 so it consults the mask (not the causal flag).
+    cpu_paged_prefill_ref(h_q_f, h_k_f, h_v_f, h_rtt, h_rpi, h_kvi, h_qoi,
+                             h_mask, q_len, q_len,
+                             B, Hq, Hkv, HEAD_DIM, max_ctx, 0, h_o_ref);
+
+    std::vector<int> q_lens(B, q_len);
+    int *d_qtb, *d_qti;
+    int num_q_tiles = make_q_tile_mapping(q_lens, &d_qtb, &d_qti);
+
+    AttentionParams<bf16> p;
+    p.batch = B; p.q_head = Hq; p.kv_head = Hkv;
+    p.head_dim = HEAD_DIM;
+    p.q_l_stride = Hq * HEAD_DIM; p.q_h_stride = HEAD_DIM; p.q_d_stride = 1;
+    p.max_context_len = max_ctx;
+    p.q_len = B * q_len;
+    p.causal_offset = -1; p.use_mask = 1;
+    p.mask = d_mask; p.mask_b_stride = q_len * q_len;
+    p.mask_h_stride = 0; p.mask_l_stride = q_len;
+    p.scale = 1.0f / sqrtf((float)HEAD_DIM);
+    p.q_ptr = d_q; p.k_ptr = d_k_pool; p.v_ptr = d_v_pool;
+    p.req_to_token = d_rtt; p.req_pool_indices = d_rpi;
+    p.kv_indptr = d_kvi; p.qo_indptr = d_qoi;
+    p.q_tile_to_batch = d_qtb; p.q_tile_to_index = d_qti;
+    p.num_q_tiles = num_q_tiles;
+    p.o_ptr = d_o; p.o_part = nullptr; p.ml_part = nullptr;
+
+    dispatch_by_head_dim(HEAD_DIM, PagedPrefillDispatch{p});
+    cudaDeviceSynchronize();
+
+    bf16* h_o_bf = (bf16*)malloc(sz_q);
+    cudaMemcpy(h_o_bf, d_o, sz_q, cudaMemcpyDeviceToHost);
+    float* h_o_got = (float*)malloc(total_q * Hq * HEAD_DIM * sizeof(float));
+    for (int i = 0; i < total_q * Hq * HEAD_DIM; i++) h_o_got[i] = bf2f(h_o_bf[i]);
+
+    const float atol = 0.01f, rtol = 0.01f;
+    bool pass = true;
+    float max_err = 0.0f;
+    for (int i = 0; i < total_q * Hq * HEAD_DIM; i++) {
+        float e = fabsf(h_o_got[i] - h_o_ref[i]);
+        if (e > max_err) max_err = e;
+        if (e > atol + rtol * fabsf(h_o_ref[i])) { pass = false; break; }
+    }
+
+    print_paged_row(cfg, max_err, pass);
+
+    free(h_q); free(h_k_pool); free(h_v_pool); free(h_rtt); free(h_rpi);
+    free(h_kvi); free(h_qoi); free(h_mask); free(h_q_f); free(h_k_f); free(h_v_f);
+    free(h_o_ref); free(h_o_bf); free(h_o_got);
+    cudaFree(d_q); cudaFree(d_o); cudaFree(d_k_pool); cudaFree(d_v_pool);
+    cudaFree(d_rtt); cudaFree(d_rpi); cudaFree(d_kvi); cudaFree(d_qoi);
+    cudaFree(d_mask);
+    cudaFree(d_qtb); cudaFree(d_qti);
+    return pass ? 0 : 1;
+}
+
+// ======================================================================
+// BENCH
+// ======================================================================
+template <int HEAD_DIM>
+static void bench_decode(int B, int Hq, int Hkv, int seq_len) {
+    int max_ctx = max(16384, seq_len + 16);
+    int pool_size = B * (seq_len + 16);
+    int num_reqs = B;
+
+    size_t sz_q  = (size_t)B * Hq * HEAD_DIM * sizeof(bf16);
+    size_t sz_kv = (size_t)pool_size * Hkv * HEAD_DIM * sizeof(bf16);
+    size_t sz_rtt = (size_t)num_reqs * max_ctx * sizeof(int);
+    size_t sz_rpi = (size_t)B * sizeof(int);
+    size_t sz_kvi = (size_t)(B + 1) * sizeof(int);
+    size_t sz_op = (size_t)B * Hq * MAX_SPLITS * HEAD_DIM * sizeof(float);
+    size_t sz_ml = (size_t)B * Hq * MAX_SPLITS * 2 * sizeof(float);
+
+    bf16 *d_q, *d_o, *d_k_pool, *d_v_pool;
+    int *d_rtt, *d_rpi;
+    int *d_kvi;
+    float *d_op, *d_ml;
+    cudaMalloc(&d_q, sz_q); cudaMalloc(&d_o, sz_q);
+    cudaMalloc(&d_k_pool, sz_kv); cudaMalloc(&d_v_pool, sz_kv);
+    cudaMalloc(&d_rtt, sz_rtt); cudaMalloc(&d_rpi, sz_rpi);
+    cudaMalloc(&d_kvi, sz_kvi);
+    cudaMalloc(&d_op, sz_op); cudaMalloc(&d_ml, sz_ml);
+
+    bf16* tmp = (bf16*)malloc(sz_kv > sz_q ? sz_kv : sz_q);
+    for (size_t i = 0; i < sz_q / sizeof(bf16); i++) tmp[i] = f2bf(randf());
+    cudaMemcpy(d_q, tmp, sz_q, cudaMemcpyHostToDevice);
+    for (size_t i = 0; i < sz_kv / sizeof(bf16); i++) tmp[i] = f2bf(randf());
+    cudaMemcpy(d_k_pool, tmp, sz_kv, cudaMemcpyHostToDevice);
+    cudaMemcpy(d_v_pool, tmp, sz_kv, cudaMemcpyHostToDevice);
+
+    int* h_rtt = (int*)malloc(sz_rtt);
+    for (int r = 0; r < num_reqs; r++)
+        for (int p = 0; p < max_ctx; p++)
+            h_rtt[r * max_ctx + p] = (r * max_ctx + p) % pool_size;
+    cudaMemcpy(d_rtt, h_rtt, sz_rtt, cudaMemcpyHostToDevice);
+    int* h_rpi = (int*)malloc(sz_rpi);
+    for (int b = 0; b < B; b++) h_rpi[b] = b;
+    cudaMemcpy(d_rpi, h_rpi, sz_rpi, cudaMemcpyHostToDevice);
+    int* h_kvi = (int*)malloc(sz_kvi);
+    h_kvi[0] = 0;
+    for (int b = 0; b < B; b++) h_kvi[b + 1] = h_kvi[b] + seq_len;
+    cudaMemcpy(d_kvi, h_kvi, sz_kvi, cudaMemcpyHostToDevice);
+
+    AttentionParams<bf16> p;
+    p.batch = B; p.q_head = Hq; p.kv_head = Hkv;
+    p.head_dim = HEAD_DIM;
+    p.q_l_stride = Hq * HEAD_DIM; p.q_h_stride = HEAD_DIM; p.q_d_stride = 1;
+    p.max_context_len = max_ctx;
+    p.causal_offset = 0; p.use_mask = 0;
+    p.mask = nullptr; p.mask_b_stride = 0;
+    p.scale = 1.0f / sqrtf((float)HEAD_DIM);
+    p.q_ptr = d_q; p.k_ptr = d_k_pool; p.v_ptr = d_v_pool;
+    p.req_to_token = d_rtt; p.req_pool_indices = d_rpi;
+    p.kv_indptr = d_kvi; p.qo_indptr = nullptr;
+    p.o_ptr = d_o; p.o_part = d_op; p.ml_part = d_ml;
+
+    auto launch = [&]() {
+        dispatch_by_head_dim(HEAD_DIM, PagedDecodeDispatch{p});
+    };
+    // Decode: q_len=1, query is the last token → attends to all [0, seq_len).
+    // FLOPs = 2 * (QK^T + PV) = 4 * B * Hq * seq_len * D.
+    double flops = 4.0 * B * Hq * (double)seq_len * HEAD_DIM;
+    BenchResult r = bench_kernel(launch, 3, 10, flops);
+
+    char cfg[64];
+    snprintf(cfg, sizeof(cfg), "DEC B=%2d Hq=%2d Hk=%d kv=%4d D=%3d",
+             B, Hq, Hkv, seq_len, HEAD_DIM);
+    print_bench_row(cfg, r);
+
+    free(tmp); free(h_rtt); free(h_rpi); free(h_kvi);
+    cudaFree(d_q); cudaFree(d_o); cudaFree(d_k_pool); cudaFree(d_v_pool);
+    cudaFree(d_rtt); cudaFree(d_rpi); cudaFree(d_kvi); cudaFree(d_op); cudaFree(d_ml);
+}
+
+template <int HEAD_DIM>
+static void bench_prefill(int B, int Hq, int Hkv, int q_len, int kv_len, int causal) {
+    int total_q = B * q_len;
+    int max_ctx = kv_len + 16;
+    int pool_size = B * max_ctx;
+    int num_reqs = B;
+
+    size_t sz_q  = (size_t)total_q * Hq * HEAD_DIM * sizeof(bf16);
+    size_t sz_kv = (size_t)pool_size * Hkv * HEAD_DIM * sizeof(bf16);
+    size_t sz_rtt = (size_t)num_reqs * max_ctx * sizeof(int);
+    size_t sz_rpi = (size_t)B * sizeof(int);
+    size_t sz_kvi = (size_t)(B + 1) * sizeof(int);
+    size_t sz_qoi = (size_t)(B + 1) * sizeof(int);
+
+    bf16 *d_q, *d_o, *d_k_pool, *d_v_pool;
+    int *d_rtt, *d_rpi;
+    int *d_kvi, *d_qoi;
+    cudaMalloc(&d_q, sz_q); cudaMalloc(&d_o, sz_q);
+    cudaMalloc(&d_k_pool, sz_kv); cudaMalloc(&d_v_pool, sz_kv);
+    cudaMalloc(&d_rtt, sz_rtt); cudaMalloc(&d_rpi, sz_rpi);
+    cudaMalloc(&d_kvi, sz_kvi); cudaMalloc(&d_qoi, sz_qoi);
+
+    bf16* tmp = (bf16*)malloc(sz_kv > sz_q ? sz_kv : sz_q);
+    for (size_t i = 0; i < sz_q / sizeof(bf16); i++) tmp[i] = f2bf(randf());
+    cudaMemcpy(d_q, tmp, sz_q, cudaMemcpyHostToDevice);
+    for (size_t i = 0; i < sz_kv / sizeof(bf16); i++) tmp[i] = f2bf(randf());
+    cudaMemcpy(d_k_pool, tmp, sz_kv, cudaMemcpyHostToDevice);
+    cudaMemcpy(d_v_pool, tmp, sz_kv, cudaMemcpyHostToDevice);
+
+    int* h_rtt = (int*)malloc(sz_rtt);
+    for (int r = 0; r < num_reqs; r++)
+        for (int p = 0; p < max_ctx; p++)
+            h_rtt[r * max_ctx + p] = (r * max_ctx + p) % pool_size;
+    cudaMemcpy(d_rtt, h_rtt, sz_rtt, cudaMemcpyHostToDevice);
+    int* h_rpi = (int*)malloc(sz_rpi);
+    for (int b = 0; b < B; b++) h_rpi[b] = b;
+    cudaMemcpy(d_rpi, h_rpi, sz_rpi, cudaMemcpyHostToDevice);
+    int* h_kvi = (int*)malloc(sz_kvi);
+    h_kvi[0] = 0;
+    for (int b = 0; b < B; b++) h_kvi[b + 1] = h_kvi[b] + kv_len;
+    cudaMemcpy(d_kvi, h_kvi, sz_kvi, cudaMemcpyHostToDevice);
+    int* h_qoi = (int*)malloc(sz_qoi);
+    h_qoi[0] = 0;
+    for (int b = 0; b < B; b++) h_qoi[b + 1] = h_qoi[b] + q_len;
+    cudaMemcpy(d_qoi, h_qoi, sz_qoi, cudaMemcpyHostToDevice);
+
+    std::vector<int> q_lens(B, q_len);
+    int *d_qtb, *d_qti;
+    int num_q_tiles = make_q_tile_mapping(q_lens, &d_qtb, &d_qti);
+
+    AttentionParams<bf16> p;
+    p.batch = B; p.q_head = Hq; p.kv_head = Hkv;
+    p.head_dim = HEAD_DIM;
+    p.q_l_stride = Hq * HEAD_DIM; p.q_h_stride = HEAD_DIM; p.q_d_stride = 1;
+    p.max_context_len = max_ctx;
+    p.q_len = B * q_len;
+    p.causal_offset = causal ? 0 : -1; p.use_mask = 0;
+    p.mask = nullptr; p.mask_b_stride = 0;
+    p.scale = 1.0f / sqrtf((float)HEAD_DIM);
+    p.q_ptr = d_q; p.k_ptr = d_k_pool; p.v_ptr = d_v_pool;
+    p.req_to_token = d_rtt; p.req_pool_indices = d_rpi;
+    p.kv_indptr = d_kvi; p.qo_indptr = d_qoi;
+    p.q_tile_to_batch = d_qtb; p.q_tile_to_index = d_qti;
+    p.num_q_tiles = num_q_tiles;
+    p.o_ptr = d_o; p.o_part = nullptr; p.ml_part = nullptr;
+
+    auto launch = [&]() {
+        dispatch_by_head_dim(HEAD_DIM, PagedPrefillDispatch{p});
+    };
+    // FLOPs = 2 * (QK^T + PV) = 4 * effective_qk_pairs * Hq * D.
+    // Non-causal: effective = q_len * kv_len.
+    // Causal: Q row qi attends to [0, causal_off + qi + 1) where
+    //   causal_off = kv_len - q_len.  Total KV accesses per request:
+    //   sum_{qi=0}^{q_len-1} (kv_len - q_len + qi + 1)
+    //   = q_len * (kv_len - q_len) + q_len * (q_len + 1) / 2.
+    double eff_kv;
+    if (causal) {
+        eff_kv = (double)q_len * (kv_len - q_len)
+               + (double)q_len * (q_len + 1) / 2.0;
+    } else {
+        eff_kv = (double)q_len * kv_len;
+    }
+    double flops = 4.0 * B * Hq * eff_kv * HEAD_DIM;
+    BenchResult r = bench_kernel(launch, 3, 10, flops);
+
+    char cfg[80];
+    snprintf(cfg, sizeof(cfg), "PRE B=%d Hq=%2d Hk=%d q=%4d kv=%4d D=%3d c=%d",
+             B, Hq, Hkv, q_len, kv_len, HEAD_DIM, causal);
+    print_bench_row(cfg, r);
+
+    free(tmp); free(h_rtt); free(h_rpi); free(h_kvi); free(h_qoi);
+    cudaFree(d_q); cudaFree(d_o); cudaFree(d_k_pool); cudaFree(d_v_pool);
+    cudaFree(d_rtt); cudaFree(d_rpi); cudaFree(d_kvi); cudaFree(d_qoi);
+    cudaFree(d_qtb); cudaFree(d_qti);
+}
+
+int main() {
+    int fail = 0;
+
+    // ===== DECODE TESTS =====
+    printf("=== Paged Decode Tests ===\n");
+    print_paged_header();
+    fail += run_decode_test<128>(1, 32, 4, 512, 0, 1);
+    fail += run_decode_test<128>(1, 32, 4, 1024, 0, 2);
+    fail += run_decode_test<128>(4, 32, 4, 512, 0, 3);
+    fail += run_decode_test<128>(8, 32, 4, 1024, 0, 4);
+    fail += run_decode_test<128>(4, 32, 8, 2048, 0, 5);
+    fail += run_decode_test<128>(1, 16, 1, 256, 0, 6);
+    fail += run_decode_test<128>(2, 8, 2, 512, 1, 7);
+    fail += run_decode_test<64>(1, 4, 2, 256, 0, 8);
+    fail += run_decode_test<256>(1, 2, 1, 256, 0, 9);
+    fail += run_decode_test<128>(16, 32, 4, 2048, 0, 10);
+    fail += run_decode_test<128>(32, 32, 4, 1024, 0, 11);
+    // Production keeps a fixed 32768-wide request table.  This forces 32
+    // splits, so seq_len > 512 gives each split multiple cp.async tiles.
+    fail += run_decode_test<64>(1, 24, 4, 1100, 0, 12, 32768, 1100);
+
+    // Decode with 2D mask (regression: mixed seq_lens + HasMask)
+    fail += run_decode_mask_test<128>(2, 8, 2, 256, 30);
+    fail += run_decode_mask_test<128>(4, 32, 4, 512, 31);
+    fail += run_decode_mask_test<64>(2, 4, 2, 128, 32);
+
+    if (fail) { printf("\nFAILED decode tests\n"); return fail; }
+
+    // ===== PREFILL TESTS =====
+    printf("\n=== Paged Prefill Tests ===\n");
+    print_paged_header();
+    // Single request, pure prefill (q_len == kv_len)
+    {
+        std::vector<int> ql = {512};
+        std::vector<int> kl = {512};
+        fail += run_prefill_test<128>(1, 32, 4, ql, kl, 1, 20);
+    }
+    {
+        std::vector<int> ql = {1024};
+        std::vector<int> kl = {1024};
+        fail += run_prefill_test<128>(1, 32, 4, ql, kl, 1, 21);
+    }
+    {
+        std::vector<int> ql = {2048};
+        std::vector<int> kl = {2048};
+        fail += run_prefill_test<128>(1, 32, 4, ql, kl, 1, 22);
+    }
+    // Ragged batch: different q_lens and kv_lens
+    {
+        std::vector<int> ql = {128, 256, 64};
+        std::vector<int> kl = {128, 256, 64};
+        fail += run_prefill_test<128>(3, 32, 4, ql, kl, 1, 23);
+    }
+    {
+        std::vector<int> ql = {64, 128, 256, 32};
+        std::vector<int> kl = {64, 128, 256, 32};
+        fail += run_prefill_test<128>(4, 32, 4, ql, kl, 1, 24);
+    }
+    // Extend: kv_len > q_len (append to existing cache)
+    {
+        std::vector<int> ql = {64, 128};
+        std::vector<int> kl = {256, 512};
+        fail += run_prefill_test<128>(2, 32, 4, ql, kl, 1, 25);
+    }
+    // Non-causal
+    {
+        std::vector<int> ql = {256, 128};
+        std::vector<int> kl = {256, 128};
+        fail += run_prefill_test<128>(2, 32, 4, ql, kl, 0, 26);
+    }
+    // Single token (q_len=1 per request, like decode but via prefill path)
+    {
+        std::vector<int> ql = {1, 1, 1, 1};
+        std::vector<int> kl = {128, 256, 64, 512};
+        fail += run_prefill_test<128>(4, 32, 4, ql, kl, 1, 27);
+    }
+    // D=64
+    {
+        std::vector<int> ql = {128, 64};
+        std::vector<int> kl = {128, 64};
+        fail += run_prefill_test<64>(2, 4, 2, ql, kl, 1, 28);
+    }
+    // D=256
+    {
+        std::vector<int> ql = {128, 64};
+        std::vector<int> kl = {128, 64};
+        fail += run_prefill_test<256>(2, 2, 1, ql, kl, 1, 29);
+    }
+
+    // Prefill with 4D causal mask (regression: single-request mask path)
+    fail += run_prefill_mask_test<128>(32, 4, 512, 40);
+    fail += run_prefill_mask_test<128>(32, 4, 1024, 41);
+    fail += run_prefill_mask_test<64>(4, 2, 256, 42);
+
+    if (fail) { printf("\nFAILED prefill tests\n"); return fail; }
+    printf("\nAll tests passed!\n");
+
+    // ===== BENCH =====
+    printf("\n===== PAGED DECODE BENCH =====\n");
+    print_bench_header();
+    bench_decode<128>(1, 32, 4, 512);
+    bench_decode<128>(1, 32, 4, 1024);
+    bench_decode<128>(1, 32, 4, 2048);
+    bench_decode<128>(1, 32, 4, 4096);
+    bench_decode<128>(1, 32, 4, 16384);
+    bench_decode<128>(4, 32, 4, 2048);
+    bench_decode<128>(16, 32, 4, 2048);
+
+    printf("\n===== PAGED PREFILL BENCH =====\n");
+    print_bench_header();
+    bench_prefill<128>(1, 32, 4, 512, 512, 0);
+    bench_prefill<128>(1, 32, 4, 1024, 1024, 0);
+    bench_prefill<128>(1, 32, 4, 2048, 2048, 0);
+    bench_prefill<128>(1, 32, 4, 2048, 2048, 1);
+    bench_prefill<128>(4, 32, 4, 2048, 2048, 1);
+    bench_prefill<128>(1, 32, 4, 4096, 4096, 1);
+
+    return 0;
+}

csrc/tests/attn_prefill_test.cu

@@ -1,169 +0,0 @@
-/*
-Pure-C test — uses shared dispatcher.
-nvcc -I csrc -arch=sm_89 -O3 \
-    --use_fast_math --ptxas-options=-O3 --extra-device-vectorization \
-    csrc/tests/attn_prefill_test.cu -o test && ./test
-*/
-
-#include "test_utils.cuh"
-#include "../kernels/attn_dispatchers.cuh"
-
-// Warmed-up, CUDA-event timed throughput sweep over the production MMA path.
-static void bench() {
-    const int cfgs[][7] = {
-        {1,32,4,512,512,128,0},
-        {1,32,4,1024,1024,128,0},
-        {1,32,4,2048,2048,128,0},
-        {1,32,4,2048,2048,128,1},
-        {4,32,4,2048,2048,128,1},
-        {1,32,4,4096,4096,128,1},
-    };
-    int n = sizeof(cfgs)/sizeof(cfgs[0]);
-    const int WARMUP = 10, ITERS = 50;
-    printf("\n===== PREFILL BENCH (warmup=%d iters=%d) =====\n", WARMUP, ITERS);
-    printf("%-46s | %10s | %10s | %10s\n",
-           "config", "latency", "bandwidth", "throughput");
-    printf("---------------------------------------------------------------"
-           "----------------------------\n");
-
-    for (int ci = 0; ci < n; ci++) {
-        int B=cfgs[ci][0], Hq=cfgs[ci][1], Hk=cfgs[ci][2];
-        int ql=cfgs[ci][3], kl=cfgs[ci][4], D=cfgs[ci][5], causal=cfgs[ci][6];
-        size_t nQ=(size_t)B*Hq*ql*D, nKV=(size_t)B*Hk*kl*D;
-
-        bf16 *dQ,*dK,*dV,*dO,*tmp;
-        cudaMalloc(&dQ,nQ*2); cudaMalloc(&dK,nKV*2);
-        cudaMalloc(&dV,nKV*2); cudaMalloc(&dO,nQ*2);
-        size_t big = nQ>nKV?nQ:nKV; tmp=new bf16[big];
-        for (size_t i=0;i<nQ;i++)  tmp[i]=f2bf(randf());
-        cudaMemcpy(dQ,tmp,nQ*2,cudaMemcpyHostToDevice);
-        for (size_t i=0;i<nKV;i++) tmp[i]=f2bf(randf());
-        cudaMemcpy(dK,tmp,nKV*2,cudaMemcpyHostToDevice);
-        for (size_t i=0;i<nKV;i++) tmp[i]=f2bf(randf());
-        cudaMemcpy(dV,tmp,nKV*2,cudaMemcpyHostToDevice);
-
-        AttentionParams<bf16> p;
-        p.batch=B; p.q_head=Hq; p.kv_head=Hk; p.q_len=ql; p.kv_len=kl; p.head_dim=D;
-        p.use_mask=0; p.causal_offset=causal?0:-1;
-        set_default_strides(p);
-        p.scale=1.0f/sqrtf((float)D);
-        p.q=dQ; p.k=dK; p.v=dV; p.mask=nullptr; p.o=dO;
-
-        auto launch = [&]() { dispatch_by_head_dim(D, [&]<int H>() { dispatch_prefill<H>(p); }); };
-        for (int i=0;i<WARMUP;i++) launch();
-        cudaDeviceSynchronize();
-        cudaError_t err=cudaGetLastError();
-        if (err!=cudaSuccess){printf("CUDA err: %s\n",cudaGetErrorString(err));return;}
-
-        cudaEvent_t s,e; cudaEventCreate(&s); cudaEventCreate(&e);
-        cudaEventRecord(s);
-        for (int i=0;i<ITERS;i++) launch();
-        cudaEventRecord(e); cudaEventSynchronize(e);
-        float ms=0; cudaEventElapsedTime(&ms,s,e); ms/=ITERS;
-
-        double flops = 4.0*B*Hq*(double)ql*kl*D;
-        if (causal) flops *= 0.5;
-        double tflops = flops/(ms*1e-3)/1e12;
-        double bytes = 2.0 * (2.0*nQ + 2.0*nKV);
-        double gbps = bytes/(ms*1e-3)/1e9;
-
-        char cfg[64];
-        snprintf(cfg, sizeof(cfg),
-                 "B=%2d Hq=%2d Hk=%d q=%4d kv=%4d D=%3d causal=%d",
-                 B,Hq,Hk,ql,kl,D,causal);
-        printf("%-46s | %7.4f ms | %7.1f GB/s | %6.2f TFLOP/s\n",
-               cfg, ms, gbps, tflops);
-
-        cudaFree(dQ);cudaFree(dK);cudaFree(dV);cudaFree(dO);
-        delete[]tmp; cudaEventDestroy(s); cudaEventDestroy(e);
-    }
-}
-
-static int run_test(int B, int Hq, int Hk, int ql, int kl, int D, int causal) {
-    printf("=== B=%d Hq=%d Hk=%d q=%d kv=%d D=%d causal=%d ===\n",
-           B,Hq,Hk,ql,kl,D,causal);
-
-    size_t nQ = B*Hq*ql*D, nKV = B*Hk*kl*D;
-    float *hQ=new float[nQ], *hK=new float[nKV], *hV=new float[nKV];
-    for (size_t i=0;i<nQ;i++) hQ[i]=randf();
-    for (size_t i=0;i<nKV;i++){hK[i]=randf();hV[i]=randf();}
-
-    bf16 *dQ,*dK,*dV,*dO,*tmp;
-    cudaMalloc(&dQ,nQ*2); cudaMalloc(&dK,nKV*2);
-    cudaMalloc(&dV,nKV*2); cudaMalloc(&dO,nQ*2);
-    tmp=new bf16[max(nQ,nKV)];
-    for (size_t i=0;i<nQ;i++) tmp[i]=f2bf(hQ[i]);
-    cudaMemcpy(dQ,tmp,nQ*2,cudaMemcpyHostToDevice);
-    for (size_t i=0;i<nKV;i++) tmp[i]=f2bf(hK[i]);
-    cudaMemcpy(dK,tmp,nKV*2,cudaMemcpyHostToDevice);
-    for (size_t i=0;i<nKV;i++) tmp[i]=f2bf(hV[i]);
-    cudaMemcpy(dV,tmp,nKV*2,cudaMemcpyHostToDevice);
-
-    AttentionParams<bf16> p;
-    p.batch=B; p.q_head=Hq; p.kv_head=Hk; p.q_len=ql; p.kv_len=kl; p.head_dim=D;
-    p.use_mask=0; p.causal_offset=causal?0:-1;
-    set_default_strides(p);
-    p.scale=1.0f/sqrtf((float)D);
-    p.q=dQ; p.k=dK; p.v=dV; p.mask=nullptr; p.o=dO;
-
-    double t0=now_ms();
-    dispatch_by_head_dim(D, [&]<int H>() { dispatch_prefill<H>(p); });
-    cudaDeviceSynchronize();
-    double kms=now_ms()-t0;
-    cudaError_t err=cudaGetLastError();
-    if (err!=cudaSuccess){printf("CUDA err: %s\n",cudaGetErrorString(err));return 1;}
-
-    bf16* hOut=new bf16[nQ];
-    cudaMemcpy(hOut,dO,nQ*2,cudaMemcpyDeviceToHost);
-
-    float* ref=new float[nQ];
-    cpu_attention_ref(hQ, hK, hV, nullptr, ref, B, Hq, Hk, ql, kl, D, causal ? 0 : -1);
-
-    float max_abs_err=0, max_rel_err=0;
-    for (size_t i=0;i<nQ;i++) {
-        float err=fabsf(bf2f(hOut[i])-ref[i]);
-        if(err>max_abs_err) max_abs_err=err;
-        float rel=err/fmaxf(fabsf(ref[i]), 1e-8f);
-        if(rel>max_rel_err) max_rel_err=rel;
-    }
-    const float atol=0.01f, rtol=0.01f;
-    bool pass=true;
-    for (size_t i=0;i<nQ;i++) {
-        float err=fabsf(bf2f(hOut[i])-ref[i]);
-        if (err > atol + rtol * fabsf(ref[i])) { pass=false; break; }
-    }
-    printf("kernel: %.3f ms  max_abs_err: %.6e  max_rel_err: %.6e  %s\n\n",
-           kms, max_abs_err, max_rel_err, pass?"PASS":"FAIL");
-
-    cudaFree(dQ);cudaFree(dK);cudaFree(dV);cudaFree(dO);
-    delete[]hQ;delete[]hK;delete[]hV;delete[]hOut;delete[]ref;delete[]tmp;
-
-    return pass ? 0 : 1;
-}
-
-int main() {
-    const int configs[][7] = {
-        {1,2,1,64,128,64,0},     // tiny: B,Hq,Hk,q,kv,D,causal
-        {1,32,4,512,512,128,0},  // standard
-        {1,32,4,128,256,128,0},  // medium
-        {1,4,2,256,256,128,1},   // causal
-    };
-    int n_configs = sizeof(configs) / sizeof(configs[0]);
-    int fail = 0;
-
-    for (int ci = 0; ci < n_configs; ci++) {
-        int B=configs[ci][0], Hq=configs[ci][1], Hk=configs[ci][2];
-        int ql=configs[ci][3], kl=configs[ci][4], D=configs[ci][5];
-        int causal=configs[ci][6];
-        fail += run_test(B, Hq, Hk, ql, kl, D, causal);
-        if (fail) break;
-    }
-
-    if (fail) {
-        printf("FAILED\n");
-        return fail;
-    }
-    printf("All tests passed!\n");
-    bench();
-    return 0;
-}

csrc/tests/attn_test.cu

@@ -0,0 +1,357 @@
+/*
+Pure-C test — uses shared dispatcher.  Combines the decode (split-KV) and
+prefill (split-Q) correctness checks + benchmarks into one binary.
+nvcc -I csrc -arch=sm_89 -O3 \
+    --use_fast_math --ptxas-options=-O3 --extra-device-vectorization \
+    -Xcompiler -fopenmp csrc/tests/attn_test.cu -o test && ./test
+*/
+
+#include "test_utils.cuh"
+#include "../kernels/attn_dispatchers.cuh"
+
+struct DecodeDispatch { AttentionParams<bf16>& p; template<int H> void operator()() { dispatch_decode<H>(p, 0); } };
+struct PrefillDispatch { AttentionParams<bf16>& p; template<int H> void operator()() { dispatch_prefill<H>(p, 0); } };
+
+// Split-K scratch (torch-free)
+struct DecodeScratch {
+    float* o_part = nullptr;
+    float* ml_part = nullptr;
+};
+
+static void setup_scratch(AttentionParams<bf16>& p, DecodeScratch& sc) {
+    int max_splits = 32;
+    cudaMalloc(&sc.o_part, (size_t)p.batch * p.q_head * max_splits * p.head_dim * sizeof(float));
+    cudaMalloc(&sc.ml_part, (size_t)p.batch * p.q_head * max_splits * 2 * sizeof(float));
+}
+
+static void free_scratch(DecodeScratch& sc) {
+    cudaFree(sc.o_part); cudaFree(sc.ml_part);
+}
+
+// ======================================================================
+// DECODE
+// ======================================================================
+
+static int run_decode_test(int B, int Hq, int Hk, int sl, int D, int causal) {
+    size_t nQ = B*Hq*1*D, nKV = B*Hk*sl*D;
+    float *hQ=new float[nQ], *hK=new float[nKV], *hV=new float[nKV];
+    for (size_t i=0;i<nQ;i++) hQ[i]=randf();
+    for (size_t i=0;i<nKV;i++){hK[i]=randf();hV[i]=randf();}
+
+    bool* hMask=new bool[B*sl];
+    for (int i=0;i<B*sl;i++) hMask[i]=true;
+
+    bf16 *dQ,*dK,*dV,*dO,*tmp;
+    bool* dMask;
+    cudaMalloc(&dQ,nQ*2); cudaMalloc(&dK,nKV*2);
+    cudaMalloc(&dV,nKV*2); cudaMalloc(&dO,nQ*2);
+    cudaMalloc(&dMask,B*sl);
+
+    tmp=new bf16[max(nQ,nKV)];
+    for (size_t i=0;i<nQ;i++) tmp[i]=f2bf(hQ[i]);
+    cudaMemcpy(dQ,tmp,nQ*2,cudaMemcpyHostToDevice);
+    for (size_t i=0;i<nKV;i++) tmp[i]=f2bf(hK[i]);
+    cudaMemcpy(dK,tmp,nKV*2,cudaMemcpyHostToDevice);
+    for (size_t i=0;i<nKV;i++) tmp[i]=f2bf(hV[i]);
+    cudaMemcpy(dV,tmp,nKV*2,cudaMemcpyHostToDevice);
+    cudaMemcpy(dMask,hMask,B*sl,cudaMemcpyHostToDevice);
+
+    AttentionParams<bf16> p;
+    p.batch=B; p.q_head=Hq; p.kv_head=Hk; p.q_len=1; p.kv_len=sl; p.head_dim=D;
+    p.use_mask=0; p.causal_offset=causal?0:-1;
+    p.scale=1.0f/sqrtf((float)D);
+    set_default_strides(p);
+    p.q_ptr=dQ; p.k_ptr=dK; p.v_ptr=dV; p.mask=nullptr; p.o_ptr=dO;
+
+    DecodeScratch sc;
+    setup_scratch(p, sc);
+    p.o_part = sc.o_part; p.ml_part = sc.ml_part;
+
+    double t0=now_ms();
+    dispatch_by_head_dim(D, DecodeDispatch{p});
+    cudaDeviceSynchronize();
+    (void)t0;
+    cudaError_t err=cudaGetLastError();
+    if (err!=cudaSuccess){printf("CUDA err: %s\n",cudaGetErrorString(err));return 1;}
+
+    bf16* hOut=new bf16[nQ];
+    cudaMemcpy(hOut,dO,nQ*2,cudaMemcpyDeviceToHost);
+
+    float* ref=new float[nQ];
+    cpu_attention_ref(hQ, hK, hV, hMask, ref, B, Hq, Hk, 1, sl, D, causal ? 0 : -1);
+
+    float max_abs_err=0, max_rel_err=0;
+    for (size_t i=0;i<nQ;i++){
+        float err=fabsf(bf2f(hOut[i])-ref[i]);
+        if(err>max_abs_err) max_abs_err=err;
+        float rel=err/fmaxf(fabsf(ref[i]), 1e-4f);
+        if(rel>max_rel_err) max_rel_err=rel;
+    }
+    const float atol=0.01f, rtol=0.01f;
+    bool pass=true;
+    for (size_t i=0;i<nQ;i++){
+        float err=fabsf(bf2f(hOut[i])-ref[i]);
+        if (err > atol + rtol * fabsf(ref[i])) { pass=false; break; }
+    }
+    char cfg[64];
+    snprintf(cfg, sizeof(cfg), "B=%2d Hq=%2d Hk=%d seq=%4d D=%3d causal=%d",
+             B, Hq, Hk, sl, D, causal);
+    print_test_row(cfg, max_abs_err, max_rel_err, pass);
+
+    cudaFree(dQ);cudaFree(dK);cudaFree(dV);cudaFree(dO);cudaFree(dMask);
+    free_scratch(sc);
+    delete[]hQ;delete[]hK;delete[]hV;delete[]hMask;delete[]hOut;delete[]ref;delete[]tmp;
+
+    return pass ? 0 : 1;
+}
+
+static void bench_decode() {
+    const int cfgs[][5] = {
+        {1, 32, 4, 512, 128},
+        {1, 32, 4, 1024, 128},
+        {1, 32, 4, 2048, 128},
+        {1, 32, 4, 4096, 128},
+        {16, 32, 4, 2048, 128},
+        {32, 32, 4, 1024, 128},
+    };
+    const int WARMUP = 3, ITERS = 10;
+    printf("\n===== DECODE BENCH (warmup=%d iters=%d) =====\n", WARMUP, ITERS);
+    print_bench_header();
+
+    int n = sizeof(cfgs) / sizeof(cfgs[0]);
+    for (int ci = 0; ci < n; ci++) {
+        int B = cfgs[ci][0], Hq = cfgs[ci][1], Hk = cfgs[ci][2];
+        int sl = cfgs[ci][3], D = cfgs[ci][4];
+        size_t nQ = (size_t)B * Hq * D;
+        size_t nKV = (size_t)B * Hk * sl * D;
+
+        bf16 *dQ, *dK, *dV, *dO;
+        cudaMalloc(&dQ, nQ*2); cudaMalloc(&dK, nKV*2);
+        cudaMalloc(&dV, nKV*2); cudaMalloc(&dO, nQ*2);
+        size_t big = nQ > nKV ? nQ : nKV; bf16* tmp = new bf16[big];
+        for (size_t i = 0; i < nQ; i++)  tmp[i] = f2bf(randf());
+        cudaMemcpy(dQ, tmp, nQ*2, cudaMemcpyHostToDevice);
+        for (size_t i = 0; i < nKV; i++) tmp[i] = f2bf(randf());
+        cudaMemcpy(dK, tmp, nKV*2, cudaMemcpyHostToDevice);
+        for (size_t i = 0; i < nKV; i++) tmp[i] = f2bf(randf());
+        cudaMemcpy(dV, tmp, nKV*2, cudaMemcpyHostToDevice);
+        delete[] tmp;
+
+        AttentionParams<bf16> p;
+        p.batch = B; p.q_head = Hq; p.kv_head = Hk; p.q_len = 1; p.kv_len = sl;
+        p.head_dim = D; p.use_mask = 0; p.causal_offset = -1;
+        p.scale = 1.0f / sqrtf((float)D);
+        set_default_strides(p);
+        p.q_ptr = dQ; p.k_ptr = dK; p.v_ptr = dV; p.mask = nullptr; p.o_ptr = dO;
+
+        DecodeScratch sc;
+        setup_scratch(p, sc);
+        p.o_part = sc.o_part; p.ml_part = sc.ml_part;
+
+        auto launch = [&]() { dispatch_by_head_dim(D, DecodeDispatch{p}); };
+        double flops = 4.0 * B * Hq * (double)sl * D;
+        BenchResult r = bench_kernel(launch, WARMUP, ITERS, flops);
+
+        char cfg[64];
+        snprintf(cfg, sizeof(cfg),
+                 "B=%2d Hq=%2d Hk=%d q=%4d kv=%4d D=%3d causal=%d",
+                 B, Hq, Hk, 1, sl, D, 0);
+        print_bench_row(cfg, r);
+
+        cudaFree(dQ); cudaFree(dK); cudaFree(dV); cudaFree(dO);
+        free_scratch(sc);
+    }
+}
+
+// ======================================================================
+// PREFILL
+// ======================================================================
+
+static int run_prefill_test(int B, int Hq, int Hk, int ql, int kl, int D, int causal) {
+    size_t nQ = B*Hq*ql*D, nKV = B*Hk*kl*D;
+    float *hQ=new float[nQ], *hK=new float[nKV], *hV=new float[nKV];
+    for (size_t i=0;i<nQ;i++) hQ[i]=randf();
+    for (size_t i=0;i<nKV;i++){hK[i]=randf();hV[i]=randf();}
+
+    bf16 *dQ,*dK,*dV,*dO,*tmp;
+    cudaMalloc(&dQ,nQ*2); cudaMalloc(&dK,nKV*2);
+    cudaMalloc(&dV,nKV*2); cudaMalloc(&dO,nQ*2);
+    tmp=new bf16[max(nQ,nKV)];
+    for (size_t i=0;i<nQ;i++) tmp[i]=f2bf(hQ[i]);
+    cudaMemcpy(dQ,tmp,nQ*2,cudaMemcpyHostToDevice);
+    for (size_t i=0;i<nKV;i++) tmp[i]=f2bf(hK[i]);
+    cudaMemcpy(dK,tmp,nKV*2,cudaMemcpyHostToDevice);
+    for (size_t i=0;i<nKV;i++) tmp[i]=f2bf(hV[i]);
+    cudaMemcpy(dV,tmp,nKV*2,cudaMemcpyHostToDevice);
+
+    AttentionParams<bf16> p;
+    p.batch=B; p.q_head=Hq; p.kv_head=Hk; p.q_len=ql; p.kv_len=kl; p.head_dim=D;
+    p.use_mask=0; p.causal_offset=causal?0:-1;
+    set_default_strides(p);
+    p.scale=1.0f/sqrtf((float)D);
+    p.q_ptr=dQ; p.k_ptr=dK; p.v_ptr=dV; p.mask=nullptr; p.o_ptr=dO;
+
+    double t0=now_ms();
+    dispatch_by_head_dim(D, PrefillDispatch{p});
+    cudaDeviceSynchronize();
+    (void)t0;
+    cudaError_t err=cudaGetLastError();
+    if (err!=cudaSuccess){printf("CUDA err: %s\n",cudaGetErrorString(err));return 1;}
+
+    bf16* hOut=new bf16[nQ];
+    cudaMemcpy(hOut,dO,nQ*2,cudaMemcpyDeviceToHost);
+
+    float* ref=new float[nQ];
+    cpu_attention_ref(hQ, hK, hV, nullptr, ref, B, Hq, Hk, ql, kl, D, causal ? 0 : -1);
+
+    float max_abs_err=0, max_rel_err=0;
+    for (size_t i=0;i<nQ;i++) {
+        float err=fabsf(bf2f(hOut[i])-ref[i]);
+        if(err>max_abs_err) max_abs_err=err;
+        float rel=err/fmaxf(fabsf(ref[i]), 1e-4f);
+        if(rel>max_rel_err) max_rel_err=rel;
+    }
+    const float atol=0.01f, rtol=0.01f;
+    bool pass=true;
+    for (size_t i=0;i<nQ;i++) {
+        float err=fabsf(bf2f(hOut[i])-ref[i]);
+        if (err > atol + rtol * fabsf(ref[i])) { pass=false; break; }
+    }
+    char cfg[64];
+    snprintf(cfg, sizeof(cfg), "B=%2d Hq=%2d Hk=%d q=%4d kv=%4d D=%3d causal=%d",
+             B, Hq, Hk, ql, kl, D, causal);
+    print_test_row(cfg, max_abs_err, max_rel_err, pass);
+
+    cudaFree(dQ);cudaFree(dK);cudaFree(dV);cudaFree(dO);
+    delete[]hQ;delete[]hK;delete[]hV;delete[]hOut;delete[]ref;delete[]tmp;
+
+    return pass ? 0 : 1;
+}
+
+static void bench_prefill() {
+    const int cfgs[][7] = {
+        {1,32,4,1024,1024,32,0},
+        {1,32,4,1024,1024,32,1},
+        {1,32,4,4096,4096,32,1},
+        {1,32,4,1024,1024,64,0},
+        {1,32,4,1024,1024,64,1},
+        {1,32,4,4096,4096,64,1},
+        {1,32,4,512,512,128,0},
+        {1,32,4,1024,1024,128,0},
+        {1,32,4,2048,2048,128,0},
+        {1,32,4,2048,2048,128,1},
+        {4,32,4,2048,2048,128,1},
+        {1,32,4,4096,4096,128,1},
+    };
+    int n = sizeof(cfgs)/sizeof(cfgs[0]);
+    const int WARMUP = 3, ITERS = 10;
+    printf("\n===== PREFILL BENCH (warmup=%d iters=%d) =====\n", WARMUP, ITERS);
+    print_bench_header();
+
+    for (int ci = 0; ci < n; ci++) {
+        int B=cfgs[ci][0], Hq=cfgs[ci][1], Hk=cfgs[ci][2];
+        int ql=cfgs[ci][3], kl=cfgs[ci][4], D=cfgs[ci][5], causal=cfgs[ci][6];
+        size_t nQ=(size_t)B*Hq*ql*D, nKV=(size_t)B*Hk*kl*D;
+
+        bf16 *dQ,*dK,*dV,*dO,*tmp;
+        cudaMalloc(&dQ,nQ*2); cudaMalloc(&dK,nKV*2);
+        cudaMalloc(&dV,nKV*2); cudaMalloc(&dO,nQ*2);
+        size_t big = nQ>nKV?nQ:nKV; tmp=new bf16[big];
+        for (size_t i=0;i<nQ;i++)  tmp[i]=f2bf(randf());
+        cudaMemcpy(dQ,tmp,nQ*2,cudaMemcpyHostToDevice);
+        for (size_t i=0;i<nKV;i++) tmp[i]=f2bf(randf());
+        cudaMemcpy(dK,tmp,nKV*2,cudaMemcpyHostToDevice);
+        for (size_t i=0;i<nKV;i++) tmp[i]=f2bf(randf());
+        cudaMemcpy(dV,tmp,nKV*2,cudaMemcpyHostToDevice);
+
+        AttentionParams<bf16> p;
+        p.batch=B; p.q_head=Hq; p.kv_head=Hk; p.q_len=ql; p.kv_len=kl; p.head_dim=D;
+        p.use_mask=0; p.causal_offset=causal?0:-1;
+        set_default_strides(p);
+        p.scale=1.0f/sqrtf((float)D);
+        p.q_ptr=dQ; p.k_ptr=dK; p.v_ptr=dV; p.mask=nullptr; p.o_ptr=dO;
+
+        auto launch = [&]() { dispatch_by_head_dim(D, PrefillDispatch{p}); };
+        for (int i=0;i<WARMUP;i++) launch();
+        cudaDeviceSynchronize();
+        cudaError_t err=cudaGetLastError();
+        if (err!=cudaSuccess){printf("CUDA err: %s\n",cudaGetErrorString(err));return;}
+
+        cudaEvent_t s,e; cudaEventCreate(&s); cudaEventCreate(&e);
+        cudaEventRecord(s);
+        for (int i=0;i<ITERS;i++) launch();
+        cudaEventRecord(e); cudaEventSynchronize(e);
+        float ms=0; cudaEventElapsedTime(&ms,s,e); ms/=ITERS;
+
+        double flops = 4.0*B*Hq*(double)ql*kl*D;
+        if (causal) flops *= 0.5;
+        double tflops = flops/(ms*1e-3)/1e12;
+        BenchResult r{ms, tflops};
+
+        char cfg[64];
+        snprintf(cfg, sizeof(cfg),
+                 "B=%2d Hq=%2d Hk=%d q=%4d kv=%4d D=%3d causal=%d",
+                 B,Hq,Hk,ql,kl,D,causal);
+        print_bench_row(cfg, r);
+
+        cudaFree(dQ);cudaFree(dK);cudaFree(dV);cudaFree(dO);
+        delete[]tmp; cudaEventDestroy(s); cudaEventDestroy(e);
+    }
+}
+
+// ======================================================================
+// MAIN
+// ======================================================================
+
+int main() {
+    int fail = 0;
+
+    // ---- DECODE ----
+    {
+        const int configs[][6] = {
+            {1, 2, 1, 64, 32, 0},
+            {1, 32, 4, 512, 128, 0},
+            {1, 32, 4, 1024, 128, 0},
+            {1, 32, 4, 512, 128, 1},
+        };
+        int n_cfgs = sizeof(configs) / sizeof(configs[0]);
+        printf("=== DECODE TESTS ===\n");
+        print_test_header();
+        for (int ci = 0; ci < n_cfgs; ci++) {
+            int B = configs[ci][0], Hq = configs[ci][1], Hk = configs[ci][2];
+            int sl = configs[ci][3], D = configs[ci][4], causal = configs[ci][5];
+            fail += run_decode_test(B, Hq, Hk, sl, D, causal);
+            if (fail) break;
+        }
+        if (fail) { printf("FAILED decode tests\n"); return fail; }
+        bench_decode();
+    }
+
+    // ---- PREFILL ----
+    {
+        const int configs[][7] = {
+            {1,2,1,64,128,32,0},     // scalar fallback D=32
+            {1,4,2,256,256,32,1},     // causal D=32 dispatch
+            {1,2,1,64,128,64,0},     // tiny: B,Hq,Hk,q,kv,D,causal
+            {1,4,2,256,256,64,1},     // causal D=64 dispatch
+            {1,32,4,512,512,128,0},  // standard
+            {1,32,4,128,256,128,0},  // medium
+            {1,4,2,256,256,128,1},   // causal
+        };
+        int n_configs = sizeof(configs) / sizeof(configs[0]);
+        printf("\n=== PREFILL TESTS ===\n");
+        print_test_header();
+        for (int ci = 0; ci < n_configs; ci++) {
+            int B=configs[ci][0], Hq=configs[ci][1], Hk=configs[ci][2];
+            int ql=configs[ci][3], kl=configs[ci][4], D=configs[ci][5];
+            int causal=configs[ci][6];
+            fail += run_prefill_test(B, Hq, Hk, ql, kl, D, causal);
+            if (fail) break;
+        }
+        if (fail) { printf("FAILED prefill tests\n"); return fail; }
+        bench_prefill();
+    }
+
+    printf("\nAll tests passed!\n");
+    return 0;
+}

csrc/tests/fp8_mma_test.cu

@@ -0,0 +1,155 @@
+/*
+Single-kernel BF16 -> FP8 MMA -> BF16 demo for Ada (sm_89).
+
+nvcc -I csrc -arch=sm_89 -std=c++17 -O3 --use_fast_math \
+    --ptxas-options=-O3,-v csrc/tests/fp8_mma_test.cu -o fp8_mma_test \
+    && ./fp8_mma_test
+*/
+
+#include "test_utils.cuh"
+
+#include <cuda_fp8.h>
+
+#include <algorithm>
+#include <vector>
+
+constexpr int M = 16;
+constexpr int N = 8;
+constexpr int K = 32;
+
+__device__ __forceinline__ unsigned pack_fp8x4(float x0, float x1, float x2,
+                                                float x3) {
+    __nv_fp8_e4m3 q0(x0);
+    __nv_fp8_e4m3 q1(x1);
+    __nv_fp8_e4m3 q2(x2);
+    __nv_fp8_e4m3 q3(x3);
+    return static_cast<unsigned>(q0.__x) |
+           (static_cast<unsigned>(q1.__x) << 8) |
+           (static_cast<unsigned>(q2.__x) << 16) |
+           (static_cast<unsigned>(q3.__x) << 24);
+}
+
+__device__ __forceinline__ unsigned load_quantize_fp8x4(
+    const bf16* src, float scale_inv) {
+    return pack_fp8x4(__bfloat162float(src[0]) * scale_inv,
+                      __bfloat162float(src[1]) * scale_inv,
+                      __bfloat162float(src[2]) * scale_inv,
+                      __bfloat162float(src[3]) * scale_inv);
+}
+
+__device__ __forceinline__ void mma_fp8_16832(float d[4],
+                                               const unsigned a[4],
+                                               const unsigned b[2]) {
+    asm volatile(
+        "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
+        "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};"
+        : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3])
+        : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]),
+          "r"(b[0]), "r"(b[1]));
+}
+
+__global__ void fused_bf16_fp8_mma_kernel(
+    const bf16* __restrict__ a, const bf16* __restrict__ b,
+    bf16* __restrict__ out, float scale_a, float scale_b) {
+    const int lane = threadIdx.x;
+    const int group = lane >> 2;
+    const int thread_in_group = lane & 3;
+    const int k0 = thread_in_group * 4;
+
+    // PTX m16n8k32 A fragment: two rows, two 16-column K partitions.
+    unsigned a_frag[4];
+    a_frag[0] = load_quantize_fp8x4(&a[group * K + k0], 1.0f / scale_a);
+    a_frag[1] = load_quantize_fp8x4(&a[(group + 8) * K + k0], 1.0f / scale_a);
+    a_frag[2] = load_quantize_fp8x4(&a[group * K + k0 + 16], 1.0f / scale_a);
+    a_frag[3] = load_quantize_fp8x4(&a[(group + 8) * K + k0 + 16],
+                                    1.0f / scale_a);
+
+    // B is supplied as row-major [N,K], equivalent to the col-major [K,N]
+    // operand required by the MMA instruction.
+    unsigned b_frag[2];
+    b_frag[0] = load_quantize_fp8x4(&b[group * K + k0], 1.0f / scale_b);
+    b_frag[1] = load_quantize_fp8x4(&b[group * K + k0 + 16], 1.0f / scale_b);
+
+    float acc[4] = {0.0f, 0.0f, 0.0f, 0.0f};
+    mma_fp8_16832(acc, a_frag, b_frag);
+
+    const int col = thread_in_group * 2;
+    const float output_scale = scale_a * scale_b;
+    *reinterpret_cast<__nv_bfloat162*>(&out[group * N + col]) =
+        __floats2bfloat162_rn(acc[0] * output_scale,
+                              acc[1] * output_scale);
+    *reinterpret_cast<__nv_bfloat162*>(&out[(group + 8) * N + col]) =
+        __floats2bfloat162_rn(acc[2] * output_scale,
+                              acc[3] * output_scale);
+}
+
+static float quantize_e4m3(float value) {
+    return static_cast<float>(__nv_fp8_e4m3(value));
+}
+
+int main() {
+    srand(0);
+    std::vector<float> a(M * K), b(N * K), reference(M * N, 0.0f);
+    std::vector<bf16> a_bf16(M * K), b_bf16(N * K), output(M * N);
+    for (float& value : a) value = randf() * 4.0f;
+    for (float& value : b) value = randf() * 4.0f;
+    for (int i = 0; i < M * K; ++i) {
+        a_bf16[i] = f2bf(a[i]);
+        a[i] = bf2f(a_bf16[i]);
+    }
+    for (int i = 0; i < N * K; ++i) {
+        b_bf16[i] = f2bf(b[i]);
+        b[i] = bf2f(b_bf16[i]);
+    }
+
+    const float amax = *std::max_element(
+        a.begin(), a.end(), [](float x, float y) { return fabsf(x) < fabsf(y); });
+    const float bmax = *std::max_element(
+        b.begin(), b.end(), [](float x, float y) { return fabsf(x) < fabsf(y); });
+    const float scale_a = fabsf(amax) / 448.0f;
+    const float scale_b = fabsf(bmax) / 448.0f;
+
+    for (int row = 0; row < M; ++row) {
+        for (int col = 0; col < N; ++col) {
+            float sum = 0.0f;
+            for (int k = 0; k < K; ++k) {
+                float qa = quantize_e4m3(a[row * K + k] / scale_a);
+                float qb = quantize_e4m3(b[col * K + k] / scale_b);
+                sum = fmaf(qa, qb, sum);
+            }
+            reference[row * N + col] = sum * scale_a * scale_b;
+        }
+    }
+
+    bf16 *d_a, *d_b, *d_out;
+    CUDA_CHECK(cudaMalloc(&d_a, a_bf16.size() * sizeof(bf16)));
+    CUDA_CHECK(cudaMalloc(&d_b, b_bf16.size() * sizeof(bf16)));
+    CUDA_CHECK(cudaMalloc(&d_out, output.size() * sizeof(bf16)));
+    CUDA_CHECK(cudaMemcpy(d_a, a_bf16.data(), a_bf16.size() * sizeof(bf16),
+                          cudaMemcpyHostToDevice));
+    CUDA_CHECK(cudaMemcpy(d_b, b_bf16.data(), b_bf16.size() * sizeof(bf16),
+                          cudaMemcpyHostToDevice));
+
+    fused_bf16_fp8_mma_kernel<<<1, 32>>>(d_a, d_b, d_out, scale_a, scale_b);
+    CUDA_CHECK(cudaDeviceSynchronize());
+    CUDA_CHECK(cudaMemcpy(output.data(), d_out, output.size() * sizeof(bf16),
+                          cudaMemcpyDeviceToHost));
+
+    float max_abs_error = 0.0f;
+    float max_rel_error = 0.0f;
+    for (int i = 0; i < M * N; ++i) {
+        float error = fabsf(bf2f(output[i]) - reference[i]);
+        max_abs_error = fmaxf(max_abs_error, error);
+        max_rel_error = fmaxf(max_rel_error,
+                              error / fmaxf(fabsf(reference[i]), 1e-4f));
+    }
+    const bool pass = max_abs_error < 0.05f;
+    print_test_header();
+    print_test_row("M=16 N=8 K=32 fused BF16->E4M3 MMA", max_abs_error,
+                   max_rel_error, pass);
+
+    cudaFree(d_a);
+    cudaFree(d_b);
+    cudaFree(d_out);
+    return pass ? 0 : 1;
+}

csrc/tests/test_utils.cuh

@@ -29,19 +29,18 @@ inline double now_ms() {
 
 struct BenchResult {
     float ms;
-    double gbps;
     double tflops;
 };
 
 template <typename Fn>
 BenchResult bench_kernel(Fn launch, int warmup, int iters,
-                         double flops, double bytes) {
+                         double flops) {
     for (int i = 0; i < warmup; i++) launch();
     cudaDeviceSynchronize();
     cudaError_t err = cudaGetLastError();
     if (err != cudaSuccess) {
         printf("CUDA error before bench: %s\n", cudaGetErrorString(err));
-        return {0, 0, 0};
+        return {0, 0};
     }
 
     cudaEvent_t s, e;
@@ -52,19 +51,33 @@ BenchResult bench_kernel(Fn launch, int warmup, int iters,
     float ms = 0; cudaEventElapsedTime(&ms, s, e); ms /= iters;
     cudaEventDestroy(s); cudaEventDestroy(e);
 
-    return {ms, bytes / (ms * 1e-3) / 1e9, flops / (ms * 1e-3) / 1e12};
+    return {ms, flops / (ms * 1e-3) / 1e12};
 }
 
 inline void print_bench_header() {
-    printf("%-46s | %10s | %10s | %10s\n",
-           "config", "latency", "bandwidth", "throughput");
+    printf("%-46s | %10s | %10s\n",
+           "config", "latency", "TFLOP/s");
     printf("---------------------------------------------------------------"
            "----------------------------\n");
 }
 
 inline void print_bench_row(const char* cfg, const BenchResult& r) {
-    printf("%-46s | %7.4f ms | %7.1f GB/s | %6.2f TFLOP/s\n",
-           cfg, r.ms, r.gbps, r.tflops);
+    printf("%-46s | %7.4f ms | %6.2f\n",
+           cfg, r.ms, r.tflops);
+}
+
+// ---- validation table (kernel vs CPU reference) ----
+inline void print_test_header() {
+    printf("%-46s | %11s | %11s | %6s\n",
+           "config", "max_abs_err", "max_rel_err", "result");
+    printf("----------------------------------------------------------------"
+           "----------------------------\n");
+}
+
+inline void print_test_row(const char* cfg, float max_abs_err,
+                           float max_rel_err, bool pass) {
+    printf("%-46s | %11.3e | %11.3e | %s\n",
+           cfg, max_abs_err, max_rel_err, pass ? "PASS" : "FAIL");
 }
 
 template <int... Ds>
@@ -94,29 +107,29 @@ void dispatch_by_head_dim(int head_dim, Fn&& fn) {
 // Set default strides for contiguous b h l d layout on AttentionParams.
 template<typename P>
 inline void set_default_strides(P& p) {
-    p.q_stride_b  = p.q_head * p.q_len * p.head_dim;
-    p.q_stride_h  = p.q_len * p.head_dim;
-    p.q_stride_l  = p.head_dim;
-    p.q_stride_d  = 1;
-    p.kv_stride_b = p.kv_head * p.kv_len * p.head_dim;
-    p.kv_stride_h = p.kv_len * p.head_dim;
-    p.kv_stride_l = p.head_dim;
-    p.kv_stride_d = 1;
+    p.q_b_stride  = p.q_head * p.q_len * p.head_dim;
+    p.q_h_stride  = p.q_len * p.head_dim;
+    p.q_l_stride  = p.head_dim;
+    p.q_d_stride  = 1;
+    p.kv_b_stride = p.kv_head * p.kv_len * p.head_dim;
+    p.kv_h_stride = p.kv_len * p.head_dim;
+    p.kv_l_stride = p.head_dim;
+    p.kv_d_stride = 1;
     p.mask_b_stride = p.kv_len;
     p.mask_h_stride = 0;
-    p.mask_q_stride = 0;
+    p.mask_l_stride = 0;
 }
 
-// Set default Q strides for contiguous b h l d layout on PagedAttentionParams.
+// Set default Q strides for a paged decode params struct.
 template<typename P>
 inline void set_default_paged_strides(P& p) {
-    p.q_stride_b  = p.q_head * p.q_len * p.head_dim;
-    p.q_stride_h  = p.q_len * p.head_dim;
-    p.q_stride_l  = p.head_dim;
-    p.q_stride_d  = 1;
+    p.q_b_stride  = p.q_head * p.q_len * p.head_dim;
+    p.q_h_stride  = p.q_len * p.head_dim;
+    p.q_l_stride  = p.head_dim;
+    p.q_d_stride  = 1;
     p.mask_b_stride = p.kv_len;
     p.mask_h_stride = 0;
-    p.mask_q_stride = 0;
+    p.mask_l_stride = 0;
 }
 
 // Generic CPU reference for multi-query / grouped-query attention.
@@ -135,9 +148,10 @@ static void cpu_attention_ref(
     float scale = 1.0f / sqrtf((float)D);
     int n_rep = Hq / Hk;
     for (int b = 0; b < B; b++) {
+        #pragma omp parallel for collapse(2) schedule(dynamic)
         for (int h = 0; h < Hq; h++) {
-            int kv_h = h / n_rep;
             for (int qi = 0; qi < q_len; qi++) {
+                int kv_h = h / n_rep;
                 float mv = -INFINITY, sv = 0.0f;
                 float accum[256] = {0.0f};
                 int lim = kv_len;

docker-compose.yml

@@ -5,7 +5,9 @@ services:
       dockerfile: Dockerfile
       args:
         CUDA_TAG: ${CUDA_TAG:-cu128}
-    user: "${UID:-1000}:${GID:-1000}"
+        USER_UID: ${ASTRAI_UID:-1000}
+        USER_GID: ${ASTRAI_GID:-1000}
+    user: "${ASTRAI_UID:-1000}:${ASTRAI_GID:-1000}"
     ports:
       - "8000:8000"
     volumes:
@@ -33,7 +35,9 @@ services:
       dockerfile: Dockerfile
       args:
         CUDA_TAG: ${CUDA_TAG:-cu128}
-    user: "${UID:-1000}:${GID:-1000}"
+        USER_UID: ${ASTRAI_UID:-1000}
+        USER_GID: ${ASTRAI_GID:-1000}
+    user: "${ASTRAI_UID:-1000}:${ASTRAI_GID:-1000}"
     ports:
       - "8000:8000"
     volumes:
@@ -46,3 +50,44 @@ services:
       retries: 3
       start_period: 120s
     restart: unless-stopped
+
+  trainer:
+    profiles: [train]
+    build:
+      context: .
+      dockerfile: Dockerfile
+      args:
+        CUDA_TAG: ${CUDA_TAG:-cu128}
+        USER_UID: ${ASTRAI_UID:-1000}
+        USER_GID: ${ASTRAI_GID:-1000}
+    init: true
+    user: "${ASTRAI_UID:-1000}:${ASTRAI_GID:-1000}"
+    volumes:
+      - ${TRAIN_DATA_DIR:-./data}:/data:ro
+      - ${TRAIN_MODEL_DIR:-./params}:/models/base:ro
+      - ${TRAIN_CHECKPOINT_DIR:-./checkpoints}:/checkpoints
+    environment:
+      - TRAIN_JOB_NAME=${TRAIN_JOB_NAME:-astrai-train}
+      - TRAIN_CONFIG=${TRAIN_CONFIG:-}
+      - BASE_MODEL=${BASE_MODEL:-/models/base}
+      - CHECKPOINT_ROOT=/checkpoints
+      - TRAIN_GPU_COUNT=${TRAIN_GPU_COUNT:-all}
+      - CUDA_VISIBLE_DEVICES
+      - NCCL_P2P_DISABLE
+      - NCCL_NET_GDR_LEVEL
+    entrypoint: ["bash", "/app/scripts/docker/train-entrypoint.sh"]
+    ipc: ${TRAIN_IPC_MODE:-host}
+    stop_grace_period: ${TRAIN_STOP_GRACE_PERIOD:-10m}
+    restart: "no"
+    logging:
+      driver: json-file
+      options:
+        max-size: ${TRAIN_LOG_MAX_SIZE:-100m}
+        max-file: ${TRAIN_LOG_MAX_FILES:-5}
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              count: all
+              capabilities: [gpu]

docs/README-zh-CN.md

@@ -14,7 +14,7 @@
 
 <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/tag/ViperEkura/AstrAI?label=Release&color=76bad9" alt="release">
   <img src="https://img.shields.io/github/stars/ViperEkura/AstrAI?style=flat&label=Stars&color=76bad9" alt="stars">
   <img src="https://img.shields.io/github/forks/ViperEkura/AstrAI?style=flat&label=Forks&color=76bad9" alt="forks">
@@ -33,7 +33,7 @@
 
 ## 📖 目录
 
-- [特性](#特性)
+- [项目概览](#项目概览)
 - [快速上手](#快速上手)
 - [演示](#演示)
 - [文档](#文档)
@@ -46,15 +46,19 @@
 <a id="chinese"></a>
 ## 中文
 
-### 特性
+### 项目概览
 
-- 🚀 **高性能**: 训练与推理双向优化，高效并行。
-- 🔧 **灵活**: 支持 seq/sft/dpo/grpo 多种训练方式，可定制模型架构。
-- 💡 **易用**: 简洁的 API 与丰富的示例、演示。
-- 📦 **轻量**: 依赖少，部署简单。
-- 🔬 **研究友好**: 模块化设计，便于实验新想法。
-- 🤗 **HuggingFace 风格 API**: 类 HuggingFace 的 AutoModel/AutoTokenizer 接口，方便加载模型和分词器。
-- 🔌 **双 API 兼容**: 同时支持 OpenAI 和 Anthropic 聊天补全 API，开箱即用。
+AstrAI 是一个覆盖模型构建、训练、评测与部署的端到端 Transformer 框架。项目以精简的 PyTorch 代码实现完整模型生命周期，包括声明式数据预处理、分布式训练、连续批处理推理，以及兼容 OpenAI 和 Anthropic 的服务接口。
+
+| 领域 | 能力 |
+|---|---|
+| **模型** | 自回归语言模型与嵌入模型，支持 GQA、MLA、MoE、RoPE，以及可扩展的 Attention/FFN 组件 |
+| **训练** | 预训练（`seq`）、监督微调（`sft`）、DPO 和 GRPO，支持梯度累积、检查点、DDP 与 FSDP |
+| **数据** | 声明式 JSON 预处理、可配置掩码与样本打包、二进制/JSONL 存储和流式数据集 |
+| **推理** | 连续批处理、分页 KV Cache、Radix 前缀缓存、流式生成，以及 Torch/CUDA/FlashAttention 后端 |
+| **服务** | 基于 FastAPI 的 OpenAI 与 Anthropic 聊天补全协议，支持 SSE 流式输出和工具调用 |
+| **评测** | Perplexity、MMLU、HumanEval、IFEval、IFD 和 ROUGE 评测工具 |
+| **扩展** | 基于工厂与注册表扩展模型、数据集、训练策略、回调、内核和协议组件 |
 
 ### 快速上手
 
@@ -62,6 +66,8 @@
 
 **1. 安装**
 
+AstrAI 需要 Python 3.12+，并精确固定 PyTorch 版本为 `2.11.0`。训练、`scripts/tools/generate.py`、生成式评估和生成演示需要 CUDA；CPU 支持仅适用于提供明确 CPU 设备路径的组件，例如 HTTP 服务和直接打分评估。
+
 ```bash
 git clone https://github.com/ViperEkura/AstrAI.git
 cd AstrAI
@@ -138,7 +144,7 @@ curl http://localhost:8000/v1/chat/completions \
 # 下载模型权重（运行演示前必需）
 python scripts/demo/download.py                      # model → params/
 
-# 交互式流式聊天（多轮对话，保持历史记录）
+# 单轮交互式流式提示循环（不保留对话历史）
 python scripts/demo/stream_chat.py
 # 在 >> 后输入消息，输入 !exit 退出
 
@@ -181,7 +187,7 @@ docker run --gpus all -it astrai:latest
 
 # 运行推理服务
 docker run --gpus all -p 8000:8000 astrai:latest \
-  python -m scripts.tools.server --port 8000 --device cuda
+  python scripts/tools/server.py --port 8000 --device cuda
 
 # 挂载数据卷
 docker run --gpus all -v /path/to/data:/data -it astrai:latest
@@ -189,7 +195,7 @@ docker run --gpus all -v /path/to/data:/data -it astrai:latest
 # Docker Compose（GPU，默认）
 docker compose up -d
 
-# Docker Compose（仅 CPU）
+# Docker Compose CPU 服务配置（不支持仅限 CUDA 的生成脚本和演示）
 docker compose --profile cpu up -d
 ```
 
@@ -239,7 +245,7 @@ SSE 流式格式、错误码和统计端点详见[推理文档](guides/inference
 
 ### 贡献
 
-我们欢迎贡献！请参阅[贡献指南](../../CONTRIBUTING.md)了解详情。
+我们欢迎贡献！请参阅[贡献指南](../CONTRIBUTING.md)了解详情。
 
 1. Fork 本仓库。
 2. 创建功能分支。
@@ -256,10 +262,10 @@ SSE 流式格式、错误码和统计端点详见[推理文档](guides/inference
 
 ### 许可证
 
-本项目采用 [GPL-3.0 许可证](../../LICENSE)。
+本项目采用 [Apache-2.0 许可证](../LICENSE)。
 
 ---
 
 <div align="center">
   <em>专为高性能与易用性设计的轻量级 Transformer 框架。</em>
-</div>
+</div>

docs/developer/architecture.md

@@ -4,7 +4,7 @@
 
 - [Class Diagram](#class-diagram) — Full Mermaid class diagram across 10+ namespaces
 - [Module Overview](#module-overview) — Component inventory per module
-- [Design Patterns](#design-patterns) — 13 documented patterns with classes
+- [Design Patterns](#design-patterns) — 15 documented patterns with classes
 - [Core Relationships](#core-relationships) — 11 key inter-component relationships
 
 ## Class Diagram
@@ -49,6 +49,11 @@ classDiagram
             +Optional[int] n_shared_experts
             +Optional[int] n_activated_experts
             +Optional[str] topk_method
+            +Optional[int] moe_intermediate_size
+            +Optional[int] shared_expert_intermediate_size
+            +bool norm_topk_prob
+            +int decoder_sparse_step
+            +Optional[List[int]] mlp_only_layers
         }
 
         class EncoderConfig {
@@ -63,6 +68,7 @@ classDiagram
             +Optional[int] num_attention_heads
             +Optional[int] num_key_value_heads
             +Optional[bool] use_qk_norm
+            +Optional[bool] use_gated_attention
             +str ffn_type
             +Optional[dict] rope_scaling
             +Optional[str] pooling_type
@@ -114,22 +120,25 @@ classDiagram
             +Dataset dataset
             +Callable optimizer_fn
             +Callable scheduler_fn
+            +Optional[str] optimizer_name
+            +Dict[str, Any] optimizer_hyperparameters
             +int n_epoch
             +int batch_per_device
             +int grad_accum_steps
             +Optional[float] max_grad_norm
             +list gradient_checkpointing_modules
+            +Optional[str] compile_mode
             +int start_epoch
             +int start_samples
             +str ckpt_dir
             +int ckpt_interval
-            +str log_dir
             +List[str] metrics
             +Optional[LoRAConfig] lora
             +int random_seed
             +int num_workers
             +Optional[int] prefetch_factor
             +bool pin_memory
+            +Optional[Callable] collate_fn
             +int nprocs
             +str backend
             +str master_addr
@@ -140,6 +149,7 @@ classDiagram
             +Optional[float] val_split
             +int val_step
             +float neftune_alpha
+            +float moe_aux_loss_coef
             +str parallel_mode
             +int rollout_interval
             +float rollout_temperature
@@ -149,7 +159,6 @@ classDiagram
             +Optional[Callable] reward_model_fn
             +dict executor_kwargs
             +dict extra_kwargs
-            +validate()
         }
 
     }
@@ -205,10 +214,6 @@ classDiagram
             -_fetch_record_key(key, index) Tensor
         }
 
-        class H5Store {
-            +load(path)
-        }
-
         class MmapStore {
             +List _mmap_refs
             +load(path)
@@ -260,11 +265,15 @@ classDiagram
     }
 
     namespace model {
-        class AutoModel {
-            +BaseModelConfig config
+        class ModelFactory {
             +Dict _entries
             +register(name) decorator
             +get_component_class(name) Type
+        }
+
+        class AutoModel {
+            <<nn.Module>>
+            +BaseModelConfig config
             +from_pretrained(path, disable_random_init, strict) nn.Module
             +save_pretrained(save_directory)
             +to(*args, **kwargs) Self
@@ -299,7 +308,14 @@ classDiagram
             +RMSNorm input_norm
             +nn.Module mlp        # MLP or DeepSeekMoE via FFNFactory
             +RMSNorm post_attention_norm
-            +forward(x, rotary_emb, attention_mask, kv_cache) Tensor
+            +forward(x, rotary_emb, attention_mask, kv_cache, is_causal) DecoderOutput
+        }
+
+        class DecoderOutput {
+            <<TypedDict>>
+            +Tensor hidden_states
+            +Optional[Tensor] aux_loss
+            +Optional[RouterStats] router_stats
         }
 
         class GQA {
@@ -314,7 +330,7 @@ classDiagram
             +Linear q_proj, k_proj, v_proj, o_proj
             +Linear gate  # only if use_gated_attention
             +RMSNorm q_norm, k_norm  # only if use_qk_norm
-            +forward(x, rotary_emb, attn_mask, kv_cache) Tensor
+            +forward(x, rotary_emb, attn_mask, kv_cache, is_causal) Tensor
         }
 
         class MLA {
@@ -334,12 +350,19 @@ classDiagram
             +Linear gate  # only if use_gated_attention
             +RMSNorm kv_norm
             +RMSNorm q_norm, k_norm  # only if use_qk_norm
-            +forward(x, rotary_emb, attn_mask, kv_cache) Tensor
+            +forward(x, rotary_emb, attn_mask, kv_cache, is_causal) Tensor
         }
 
         class MLP {
             +Linear up, gate, down
-            +forward(x) Tensor
+            +forward(x) FFNOutput
+        }
+
+        class FFNOutput {
+            <<TypedDict>>
+            +Tensor hidden_states
+            +Optional[Tensor] aux_loss
+            +Optional[RouterStats] router_stats
         }
 
         class DeepSeekMoE {
@@ -351,7 +374,7 @@ classDiagram
             +Linear router
             +ModuleList shared_experts
             +ModuleList routed_experts
-            +forward(x) Tensor
+            +forward(x) FFNOutput
         }
 
         class AttnFactory {
@@ -380,9 +403,8 @@ classDiagram
             +int max_len
             +float base
             +Optional[Dict] rope_scaling
-            +Tensor cos_table
-            +Tensor sin_table
-            +forward(x, position_ids=None) Tuple[Tensor, Tensor]
+            +Tensor freqs_cis
+            +forward(x, position_ids=None) Tensor
         }
 
         class Embedding {
@@ -486,10 +508,6 @@ classDiagram
             +save(output_dir, domain, shard_idx, tensors)
         }
 
-        class H5Writer {
-            +save(output_dir, domain, shard_idx, tensors)
-        }
-
         class Pipeline {
             +PipelineConfig config
             +List[str] paths
@@ -559,7 +577,7 @@ classDiagram
         class Trainer {
             +TrainConfig train_config
             +List[TrainCallback] callbacks
-            +train(resume_dir)
+            +train(param_path=None, resume=False)
             -_get_default_callbacks() List[TrainCallback]
         }
 
@@ -576,13 +594,17 @@ classDiagram
             +int epoch
             +int consumed_samples
             +float loss
-            +float grad_norm
+            +Dict[str, float] metrics
+            +Optional[float] grad_norm
+            +GradSNRTracker grad_snr_tracker
             +DataLoader val_dataloader
-            +float val_loss
+            +Optional[float] val_loss
             +int world_size
             +int rank
             +dict kwargs
-            +optimizer_step() int
+            +stop_requested (property) bool
+            +optimizer_step (property) int
+            +request_stop()
         }
 
         class TrainContextBuilder {
@@ -594,11 +616,22 @@ classDiagram
         class BaseStrategy {
             +Callable model
             +Optional[BaseExecutor] executor
-            +Optional[Callable] model_fn
+            +float moe_aux_loss_coef
             +dict extra_kwargs
             +str device
-            +__call__(batch) Tensor
+            +__call__(batch) LossOutput
             +compute_loss(batch) Tensor
+            +compute_loss_output(batch) LossOutput
+            +supports_online() bool
+            +set_rollout_runner(runner)
+            +prepare_from_rollout(result) Dict
+            +on_optimizer_step()
+        }
+
+        class LossOutput {
+            <<TypedDict>>
+            +Tensor loss
+            +Dict[str, float] metrics
         }
 
         class StrategyFactory {
@@ -636,9 +669,12 @@ classDiagram
 
         class RawRollout {
             +Tensor prompts
+            +Tensor prompt_mask
             +Tensor responses
             +Tensor response_mask
             +Tensor logprobs_old
+            +List[str] prompt_texts
+            +List[List[str]] response_texts
         }
 
         class RolloutResult {
@@ -647,10 +683,18 @@ classDiagram
 
         class BaseRewardModel {
             <<abstract>>
-            +score(prompts, responses) Tensor
+            +score(List[str] prompts, List[List[str]] responses) Tensor
         }
 
         class RolloutGenerator {
+            +InferenceScheduler scheduler
+            +int max_tokens
+            +int group_size
+            +float temperature
+            +int top_k
+            +float top_p
+            +float frequency_penalty
+            +int rep_window
             +generate(batch) RawRollout
         }
 
@@ -740,7 +784,7 @@ classDiagram
         }
 
         class MetricCallback {
-            +Path log_dir
+            +Path ckpt_dir
             +int save_interval
             +List[str] metrics
             +int val_step
@@ -764,9 +808,8 @@ classDiagram
             +nn.Module model
             +AutoTokenizer tokenizer
             +InferenceScheduler scheduler
-            +generate(prompt, stream, max_tokens, temperature, top_p, top_k) Union[Generator, str, List[str]]
-            +generate_with_request(request) Union[Generator, str, List[str]]
-            +generate_async(prompt, max_tokens, temperature, top_p, top_k) AsyncGenerator
+            +generate(prompt, stream, max_tokens, temperature, top_p, top_k, frequency_penalty, rep_window) Union[Generator, str, List[str]]
+            +generate_async(prompt, max_tokens, temperature, top_p, top_k, frequency_penalty, rep_window) AsyncGenerator
             +get_stats() Dict
             +shutdown()
         }
@@ -774,18 +817,37 @@ classDiagram
         class Executor {
             +AutoModel model
             +AutoTokenizer tokenizer
-            +KVCache page_cache
+            +PagePool kv_cache
+            +InferenceWorkspace _workspace
             +Optional[str] device
             +Optional[torch.dtype] dtype
-            +execute_prefill(tasks, prompt_len, start_pos)
-            +execute_decode(tasks) List[int]
+            +execute_prefill(tasks, prompt_len, start_pos=0)
+            +execute_decode(tasks, return_logprobs=False) Union[List[int], List[Tuple[int, float]]]
+        }
+
+        class InferenceWorkspace {
+            +int max_batch_size
+            +int max_seq_len
+            +torch.device device
+            +torch.dtype dtype
+            +Tensor arange
+            +Tensor input_mask
+            +Tensor input_ids
+            +Tensor req_pool_indices
+            +Tensor seq_lens
+            +Tensor kv_indptr
+            +Tensor qo_indptr
+            +Tensor inc
+            +Tensor out_cache_loc
+            +fill_input_ids(ids) Tensor
+            +decode_mask(position_ids, total_len) Tensor
         }
 
         class InferenceScheduler {
-            +KVCache _page_cache
+            +PagePool _cache
             +Executor _executor
             +TaskManager _task_mgr
-            +bool _running
+            +Event _stop_event
             +Thread _loop_thread
             +int max_seq_len
             +str device
@@ -795,6 +857,7 @@ classDiagram
             +start()
             +stop()
             +get_stats() Dict
+            +run_batch(prompt_ids_list, max_tokens, temperature, top_p, top_k, frequency_penalty, rep_window, return_logprobs) Union[List[List[int]], List[Tuple[List[int], List[float]]]]
         }
 
         class Allocator {
@@ -808,22 +871,21 @@ classDiagram
             +ref_count(idx) int
         }
 
-        class PrefixCache {
+        class RadixNode {
+            +RadixNode parent
+            +Dict children
+            +Optional[int] page_idx
+            +Tuple tokens
+            +int lock_ref
+        }
+
+        class RadixCache {
             +int _page_size
             +evict(page_idx)
             +has_page(idx) bool
             +lookup(token_ids) List[int]
             +record(page_idx, token_ids, logical_page_idx)
-        }
-
-        class PagePool {
-            -Allocator _alloc
-            -PrefixCache _prefix
-            +alloc() int
-            +free(idx)
-            +inc_ref(idx)
-            +lookup(token_ids) List[int]
-            +record(page_idx, token_ids, logical_page_idx)
+            +release(pages)
         }
 
         class KVStorage {
@@ -852,8 +914,11 @@ classDiagram
             +Tensor seq_lens
             +Tensor out_cache_loc
             +int max_len
-            +Optional[Tensor] page_table
-            +Optional[Tensor] decode_mask
+            +Optional[Tensor] kv_indptr
+            +Optional[Tensor] qo_indptr
+            +Optional[Tensor] decode_o_part
+            +Optional[Tensor] decode_ml_part
+            +Optional[Tensor] decode_out
         }
 
         class PagePool {
@@ -862,13 +927,13 @@ classDiagram
             -KVStorage _storage
             -ReqToTokenPool _req_pool
             -Allocator _alloc
-            -PrefixCache _prefix
+            -RadixCache _prefix
             +task_alloc(task_id, prompt_ids) bool
             +task_free(task_id)
             +task_extend(task_id, pos) bool
             +task_cached(task_id) int
             +task_record_hashes(task_id, prompt_ids, start_logical_page)
-            +bind_tasks(task_ids, seq_lens, device, start_pos) KVCache
+            +bind_tasks(task_ids, workspace, device, start_pos) KVCache
         }
 
     class Task {
@@ -878,6 +943,8 @@ classDiagram
         +float temperature
         +float top_p
         +int top_k
+        +float frequency_penalty
+        +int rep_window
         +TaskStatus status
         +List output_ids
         +int input_tokens
@@ -917,33 +984,24 @@ classDiagram
             +get_stats() Dict
         }
 
-        class GenerationRequest {
-            +List[Dict] messages
-            +int top_k
-            +float top_p
-            +float temperature
-            +Optional[int] max_tokens
-            +bool stream
-        }
-
         class BaseSamplingStrategy {
             <<abstract>>
-            +apply(logits, filter_value) Tensor
+            +apply(logits, filter_value, input_ids, input_mask) Tensor
         }
 
         class TemperatureStrategy {
             +float temperature
-            +apply(logits, filter_value) Tensor
+            +apply(logits, filter_value, input_ids, input_mask) Tensor
         }
 
         class TopKStrategy {
             +int top_k
-            +apply(logits, filter_value) Tensor
+            +apply(logits, filter_value, input_ids, input_mask) Tensor
         }
 
         class TopPStrategy {
             +float top_p
-            +apply(logits, filter_value) Tensor
+            +apply(logits, filter_value, input_ids, input_mask) Tensor
         }
 
         class FrequencyPenaltyStrategy {
@@ -953,8 +1011,8 @@ classDiagram
 
         class SamplingPipeline {
             +List[BaseSamplingStrategy] strategies
-            +apply(logits, filter_value) Tensor
-            +sample(logits, filter_value) Tensor
+            +apply(logits, filter_value, input_ids, input_mask) Tensor
+            +sample(logits, filter_value, input_ids, input_mask, return_logprobs) Union[Tensor, Tuple[Tensor, Tensor]]
         }
 
         class StreamDecoder {
@@ -1029,7 +1087,7 @@ classDiagram
             <<abstract>>
             +prepare(request, engine) Tuple[str, GenContext, List[str]]
             +format_stream_start(ctx) List[str]
-            +format_chunk(token) List[str]
+            +format_chunk(token, **kwargs) List[str]
             +format_stream_end(ctx, stop) List[str]
             +format_response(ctx, content, stop) Dict
         }
@@ -1037,7 +1095,7 @@ classDiagram
         class OpenAIResponseBuilder {
             +prepare(request, engine) Tuple
             +format_stream_start(ctx) List[str]
-            +format_chunk(token) List[str]
+            +format_chunk(token, **kwargs) List[str]
             +format_stream_end(ctx, stop) List[str]
             +format_response(ctx, content, stop) Dict
         }
@@ -1045,7 +1103,7 @@ classDiagram
         class AnthropicResponseBuilder {
             +prepare(request, engine) Tuple
             +format_stream_start(ctx) List[str]
-            +format_chunk(token) List[str]
+            +format_chunk(token, **kwargs) List[str]
             +format_stream_end(ctx, stop) List[str]
             +format_response(ctx, content, stop) Dict
         }
@@ -1153,10 +1211,13 @@ classDiagram
 
         class BaseExecutor {
             +GradientState gradient_state
-            +prepare(model_fn, optimizer_fn, scheduler_fn, before_wrap) tuple
+            +prepare(model_fn, optimizer_fn, scheduler_fn, before_wrap, after_wrap) tuple
             +accumulate(model) context manager
             +backward(loss)
             +unwrap_model(model) dict
+            +checkpoint_context(model) context manager
+            +clip_grad_norm(model, max_norm) float
+            +use_distributed (property) bool
             +sync_gradients (property) bool
             +grad_accum_steps (property) int
         }
@@ -1173,7 +1234,8 @@ classDiagram
         class FSDPExecutor {
             -_prepare_model(model) nn.Module
             -_no_sync(model) context manager
-            +unwrap_model(model) dict
+            +unwrap_model(model) Optional[dict]
+            +clip_grad_norm(model, max_norm) float
         }
 
         class ExecutorFactory {
@@ -1182,33 +1244,6 @@ classDiagram
             +create(parallel_mode, **kwargs) BaseExecutor
         }
 
-        class ParallelModel {
-            +dist.ProcessGroup process_group
-            +int rank
-            +int world_size
-        }
-
-        class ColumnParallelLinear {
-            +int in_features
-            +int out_features
-            +int out_features_per_rank
-            +bool gather_results
-            +Parameter weight
-            +Optional[Parameter] bias
-            +forward(x) Tensor
-            +load_state_dict(state_dict)
-        }
-
-        class RowParallelLinear {
-            +int in_features
-            +int out_features
-            +int in_features_per_rank
-            +bool reduce_results
-            +Parameter weight
-            +Optional[Parameter] bias
-            +forward(x) Tensor
-            +load_state_dict(state_dict)
-        }
     }
 
     %% Relationships — UML notation: <|-- generalization, *-- composition, o-- aggregation, --> association, ..> dependency
@@ -1230,11 +1265,8 @@ classDiagram
     BaseDataset <|-- SFTDataset
     BaseDataset <|-- DPODataset
     BaseDataset <|-- GRPODataset
-    Store <|-- H5Store
     Store <|-- MmapStore
     Store <|-- JsonlStore
-    H5Store --|> Streamable
-    H5Store --|> Recordable
     MmapStore --|> Streamable
     MmapStore --|> Recordable
     JsonlStore --|> Streamable
@@ -1243,8 +1275,6 @@ classDiagram
     BaseSamplingStrategy <|-- TopKStrategy
     BaseSamplingStrategy <|-- TopPStrategy
     BaseSamplingStrategy <|-- FrequencyPenaltyStrategy
-    ParallelModel <|-- RowParallelLinear
-    ParallelModel <|-- ColumnParallelLinear
     AutoModel <|-- AutoRegressiveLM
     AutoModel <|-- EmbeddingEncoder
     BaseConfig <|-- BaseModelConfig
@@ -1255,7 +1285,7 @@ classDiagram
     BaseConfig <|-- PipelineConfig
     BaseModelConfig <|-- AutoRegressiveLMConfig
     BaseModelConfig <|-- EncoderConfig
-    BaseFactory <|-- AutoModel
+    BaseFactory <|-- ModelFactory
     BaseFactory <|-- AttnFactory
     BaseFactory <|-- FFNFactory
     BaseFactory <|-- DatasetFactory
@@ -1286,7 +1316,6 @@ classDiagram
     PositionIdStrategy <|-- DocResetPositionId
     PositionIdStrategy <|-- ContinuousPositionId
     StoreWriter <|-- BinWriter
-    StoreWriter <|-- H5Writer
     RawRollout <|-- RolloutResult
     LaunchStrategy <|-- TorchrunStrategy
     LaunchStrategy <|-- LocalStrategy
@@ -1294,10 +1323,12 @@ classDiagram
     PagePool *-- KVStorage
     PagePool *-- ReqToTokenPool
     PagePool *-- Allocator
-    PagePool *-- PrefixCache
+    PagePool *-- RadixCache
+    RadixCache *-- RadixNode
     InferenceEngine *-- InferenceScheduler
     InferenceScheduler *-- PagePool
     InferenceScheduler *-- Executor
+    Executor *-- InferenceWorkspace
     InferenceScheduler *-- TaskManager
     AutoRegressiveLM *-- DecoderBlock
     AutoRegressiveLM *-- RotaryEmbedding
@@ -1317,8 +1348,6 @@ classDiagram
     %% --- Aggregation (weak ownership) ---
     AutoModel o-- BaseModelConfig
     AutoTokenizer o-- ChatTemplate
-    PagePool o-- Allocator
-    PagePool o-- PrefixCache
     Trainer o-- TrainCallback
     TrainContext o-- BaseStrategy
     TrainContext o-- BaseScheduler
@@ -1352,11 +1381,12 @@ classDiagram
     FFNFactory ..> DeepSeekMoE : creates
     DecoderBlock ..> AttnFactory : uses
     DecoderBlock ..> FFNFactory : uses
-    StoreFactory ..> H5Store : creates
     StoreFactory ..> MmapStore : creates
     StoreFactory ..> JsonlStore : creates
     ConfigFactory ..> AutoRegressiveLMConfig : creates
     ConfigFactory ..> EncoderConfig : creates
+    ModelFactory ..> AutoRegressiveLM : creates
+    ModelFactory ..> EmbeddingEncoder : creates
     ExecutorFactory ..> NoneExecutor : creates
     ExecutorFactory ..> DDPExecutor : creates
     ExecutorFactory ..> FSDPExecutor : creates
@@ -1369,7 +1399,8 @@ classDiagram
     Checkpoint ..> Checkpoint : serializes
     CheckpointCallback ..> Checkpoint : creates
     PagePool ..> KVCache : binds
-    InferenceEngine ..> GenerationRequest : uses
+    PagePool ..> InferenceWorkspace : fills
+    InferenceEngine ..> GenerateResult : uses
     InferenceEngine ..> GenerateResult : creates
     OpenAIResponseBuilder ..> ChatCompletionRequest : receives
     AnthropicResponseBuilder ..> MessagesRequest : receives
@@ -1399,14 +1430,14 @@ classDiagram
 | Module | Components | Description |
 |--------|------------|-------------|
 | **astrai.config** | BaseConfig, BaseModelConfig, AutoRegressiveLMConfig, EncoderConfig, ConfigFactory, TrainConfig, PipelineConfig, InputConfig, ProcessingConfig, OutputConfig | Configuration management (to_dict/from_dict, to_file/from_file) |
-| **astrai.preprocessing** | SectionRenderer, BaseMaskBuilder, MaskBuilderFactory, SectionedMaskBuilder, SingleOutputMaskBuilder, MultiOutputMaskBuilder, Pipeline, TokenizeTransform, PackingStrategy, PackingStrategyFactory, SimplePacking, BFDPacking, BFDSplitPacking, PositionIdStrategy, PositionIdStrategyFactory, NoPositionId, DocResetPositionId, ContinuousPositionId, StoreWriter, StoreWriterFactory, BinWriter, H5Writer | Declarative JSON-driven data preprocessing |
-| **astrai.dataset** | BaseDataset, SEQDataset, SFTDataset, DPODataset, GRPODataset, Store, Streamable, Recordable, H5Store, MmapStore, JsonlSource, JsonlStore, StoreFactory, RDSampler, DatasetFactory | Dataset loading and management |
+| **astrai.preprocessing** | SectionRenderer, BaseMaskBuilder, MaskBuilderFactory, SectionedMaskBuilder, SingleOutputMaskBuilder, MultiOutputMaskBuilder, Pipeline, TokenizeTransform, PackingStrategy, PackingStrategyFactory, SimplePacking, BFDPacking, BFDSplitPacking, PositionIdStrategy, PositionIdStrategyFactory, NoPositionId, DocResetPositionId, ContinuousPositionId, StoreWriter, StoreWriterFactory, BinWriter | Declarative JSON-driven data preprocessing |
+| **astrai.dataset** | BaseDataset, SEQDataset, SFTDataset, DPODataset, GRPODataset, Store, Streamable, Recordable, MmapStore, JsonlSource, JsonlStore, StoreFactory, RDSampler, DatasetFactory | Dataset loading and management |
 | **astrai.serialization** | Checkpoint | Model serialization |
-| **astrai.model** | AutoModel, AutoRegressiveLM, EmbeddingEncoder, DecoderBlock, GQA, MLA, MLP, DeepSeekMoE, AttnFactory, FFNFactory, RMSNorm, Linear, LoRAConfig, LoRALinear, RotaryEmbedding, Embedding | Neural network model |
+| **astrai.model** | ModelFactory, AutoModel, AutoRegressiveLM, EmbeddingEncoder, DecoderBlock, GQA, MLA, MLP, DeepSeekMoE, AttnFactory, FFNFactory, RMSNorm, Linear, LoRAConfig, LoRALinear, RotaryEmbedding, Embedding | Neural network model |
 | **astrai.tokenize** | AutoTokenizer, ChatTemplate | Tokenizer and chat template |
 | **astrai.trainer** | Trainer, TrainContext, TrainContextBuilder, BaseStrategy–GRPOStrategy, StrategyFactory, BaseScheduler–WSDScheduler, SchedulerFactory, TrainCallback(Protocol)–MetricCallback, CallbackFactory, RawRollout, RolloutResult, BaseRewardModel, RolloutGenerator, RolloutRunner | Training workflow |
-| **astrai.inference** | InferenceEngine, InferenceScheduler, Executor, PagePool, KVStorage, ReqToTokenPool, KVCache, Allocator, PrefixCache, Task, TaskManager, TaskStatus, StreamDecoder, GenerationRequest, GenerateResult, BaseSamplingStrategy–SamplingPipeline, FrequencyPenaltyStrategy, ProtocolHandler, ResponseBuilder, OpenAIResponseBuilder, AnthropicResponseBuilder, StopChecker, GenContext, StopInfo, ChatMessage, FunctionDef, ToolDef, ChatCompletionRequest, AnthropicMessage, MessagesRequest, BaseToolParser, ToolParserFactory, SimpleJsonToolParser | Inference service |
-| **astrai.extension** | AttentionBackend, TorchNativeBackend, CudaBackend, attn_backend, ATTN_BACKEND, attn_decode, attn_prefill, attn_paged_decode, rotary_emb, apply_rotary_emb, rotary_backend, is_available | CUDA attention + rotary kernels, backend abstraction, auto-dispatch |
+| **astrai.inference** | InferenceEngine, InferenceScheduler, Executor, InferenceWorkspace, PagePool, KVStorage, ReqToTokenPool, KVCache, Allocator, RadixCache, Task, TaskManager, TaskStatus, StreamDecoder, GenerateResult, BaseSamplingStrategy–SamplingPipeline, FrequencyPenaltyStrategy, ProtocolHandler, ResponseBuilder, OpenAIResponseBuilder, AnthropicResponseBuilder, StopChecker, GenContext, StopInfo, ChatMessage, FunctionDef, ToolDef, ChatCompletionRequest, AnthropicMessage, MessagesRequest, BaseToolParser, ToolParserFactory, SimpleJsonToolParser | Inference service |
+| **astrai.extension** | `backend` policy package, `ops` kernel-wrapper package, AttentionBackend, TorchNativeBackend, CudaBackend, FlashAttnBackend, attention, attn_backend, ATTN_BACKEND, apply_rotary_emb, is_available | Stable API over attention/rotary execution policy and optional CUDA kernels |
 | **astrai.parallel** | spawn_parallel_fn, setup_parallel, get_rank/get_world_size/get_current_device, only_on_rank, LaunchStrategy, TorchrunStrategy, LocalStrategy, BaseExecutor, ExecutorFactory, NoneExecutor, DDPExecutor, FSDPExecutor, GradientState, AccumOptimizer, AccumScheduler | Distributed parallel & gradient accumulation |
 | **astrai.factory** | BaseFactory | Component registration |
 | **astrai.protocols** | OptimizerProtocol, SchedulerProtocol | Structural subtyping for optimizer/scheduler wrappers |
@@ -1415,21 +1446,21 @@ classDiagram
 
 | Pattern | Classes | Purpose |
 |---------|---------|---------|
-| **Factory** | `AttnFactory`, `FFNFactory`, `StrategyFactory`, `DatasetFactory`, `SchedulerFactory`, `CallbackFactory`, `StoreFactory`, `ConfigFactory`, `ExecutorFactory`, `MaskBuilderFactory`, `StoreWriterFactory`, `PackingStrategyFactory`, `PositionIdStrategyFactory`, `ToolParserFactory` | Decorator-based component creation |
+| **Factory** | `ModelFactory`, `AttnFactory`, `FFNFactory`, `StrategyFactory`, `DatasetFactory`, `SchedulerFactory`, `CallbackFactory`, `StoreFactory`, `ConfigFactory`, `ExecutorFactory`, `MaskBuilderFactory`, `StoreWriterFactory`, `PackingStrategyFactory`, `PositionIdStrategyFactory`, `ToolParserFactory` | Decorator-based component creation |
 | **Registry** | `BaseFactory` | Component registration |
 | **Strategy** | `SEQStrategy`, `SFTStrategy`, `DPOStrategy`, `GRPOStrategy` | Training strategy switching |
-| **Strategy (Sampling)** | `TemperatureStrategy`, `TopKStrategy`, `TopPStrategy`, `SamplingPipeline` | Composable logit transformations |
+| **Strategy (Sampling)** | `TemperatureStrategy`, `TopKStrategy`, `TopPStrategy`, `FrequencyPenaltyStrategy`, `SamplingPipeline` | Composable logit transformations |
 | **Strategy (API)** | `ResponseBuilder`, `OpenAIResponseBuilder`, `AnthropicResponseBuilder` | HTTP API handler with format hooks |
 | **Builder** | `TrainContextBuilder` | Chain-building training context |
 | **Observer** | `TrainCallback`, callback implementations | Training process monitoring |
 | **Context** | `TrainContext` | Unified training state bag |
 | **Object Pool** | `Allocator`, `PagePool` | Page-based KV cache with LRU eviction |
-| **Strategy (Attention)** | `AttentionBackend`, `TorchNativeBackend`, `CudaBackend` | Attention computation backend switching via context manager |
-| **Auto-dispatch (Rotary)** | `apply_rotary_emb`, `rotary_backend.py`, `rotary_ops.py` | Rotary embedding CUDA kernel auto-dispatch with torch fallback |
+| **Strategy (Attention)** | `AttentionBackend`, `CudaBackend`, `FlashAttnBackend`, `TorchNativeBackend` | Attention computation backend switching via context manager |
+| **Auto-dispatch (Rotary)** | `apply_rotary_emb`, `backend/rotary.py`, `ops/rotary.py` | Rotary embedding CUDA kernel auto-dispatch with torch fallback |
 | **Executor** | `BaseExecutor`, `NoneExecutor`, `DDPExecutor`, `FSDPExecutor` | Gradient accumulation & model distribution |
-| **Storage** | `Store`, `H5Store`, `MmapStore`, `JsonlStore` | Format-agnostic data access with multi-segment support |
+| **Storage** | `Store`, `MmapStore`, `JsonlStore` | Format-agnostic data access with multi-segment support |
 | **Producer-Consumer** | `InferenceScheduler`, `Task`, queues | Continuous batching |
-| **AutoModel Registry** | `AutoModel`, `AutoRegressiveLM`, `EmbeddingEncoder` | Model-type dynamic loading |
+| **Model Registry** | `ModelFactory`, `AutoRegressiveLM`, `EmbeddingEncoder` | Model-type dynamic loading |
 
 ## Core Relationships
 
@@ -1437,12 +1468,12 @@ classDiagram
 2. **Training Flow**: `Trainer` → `TrainContextBuilder` → `TrainContext`, uses `BaseStrategy` for loss, `BaseExecutor` for gradient accumulation + model distribution
 3. **Strategy Selection**: `StrategyFactory` creates strategy by `train_type`
 4. **Executor Selection**: `ExecutorFactory.create(cfg.parallel_mode, grad_accum_steps=cfg.grad_accum_steps, **cfg.executor_kwargs)` → `NoneExecutor` / `DDPExecutor` / `FSDPExecutor`
-5. **Inference Flow**: `InferenceEngine` → `InferenceScheduler` → `AutoRegressiveLM`, backed by `PagePool` + `KVCache` + `SamplingPipeline`. Attention backend selected via `attn_backend()` context manager (`TorchNativeBackend` default, `CudaBackend` for CUDA kernels). Rotary embedding auto-dispatches to CUDA kernel when available (inference mode), else torch complex multiply (training).
+5. **Inference Flow**: `InferenceEngine` → `InferenceScheduler` → `AutoRegressiveLM`, backed by `PagePool` + `KVCache` + `SamplingPipeline`. `astrai.extension.backend` owns attention/rotary dispatch, fallback, and KV cache policy; it calls the stateless compiled-kernel wrappers in `astrai.extension.ops`. Attention uses cuda > flash > torch priority unless explicitly selected by `ASTR_BACKEND` or `attn_backend()`. Rotary embedding auto-dispatches to the CUDA op when supported, else torch complex multiply.
 6. **Distributed**: `spawn_parallel_fn` + `setup_parallel` for multi-process DDP
-7. **Dataset Loading**: `DatasetFactory` creates datasets, `Store` (H5Store/MmapStore/JsonlStore) loads data with explicit `_length` and multi-segment `_data`
-8. **Checkpoint**: `Checkpoint` saves/loads safetensors + metadata (rank-0 only), extra state saved as `{key}.pt`
+7. **Dataset Loading**: `DatasetFactory` creates datasets, `Store` (`MmapStore`/`JsonlStore`) loads data with explicit `_length` and multi-segment `_data`
+8. **Checkpoint**: `Checkpoint` saves/loads safetensors + metadata; `CheckpointCallback` performs rank-0 training saves, with extra state saved as `{key}.pt`
 9. **Scheduler**: `SchedulerFactory` creates `CosineScheduler`/`SGDRScheduler`/`WSDScheduler`
 10. **AutoModel**: `from_pretrained()` loads `config.json` + `model.safetensors`, `_disable_random_init` replaces `nn.init.*` with no-ops
 11. **Protocols**: `OptimizerProtocol` / `SchedulerProtocol` — structural subtyping for `AccumOptimizer` / `AccumScheduler` wrappers
 
-> Document Update Time: 2026-07-31
+> Document Update Time: 2026-08-16

docs/developer/cuda_kernels.md

@@ -9,6 +9,7 @@ AstrAI includes optional custom CUDA kernels for attention and rotary embedding.
 | `attn_decode` | `attn_decode.cu` | GQA decode attention (split-KV) |
 | `attn_prefill` | `attn_prefill.cu` | GQA prefill attention (split-Q) |
 | `attn_paged_decode` | `attn_paged_decode.cu` | Paged KV cache decode attention |
+| `attn_paged_prefill` | `attn_paged_prefill.cu` | Paged KV cache prefill attention (ragged batch) |
 | `rotary_emb` | `rotary_emb.cu` | Fused rotary embedding (cos/sin lookup + rotation) |
 
 Additionally, optimized `.cuh` variants with tensor-core MMA (Matrix Multiply-Accumulate) exist:
@@ -17,7 +18,11 @@ Additionally, optimized `.cuh` variants with tensor-core MMA (Matrix Multiply-Ac
 |---------|------|--------------|
 | Split-KV MMA decode | `attn_decode_split_kv_mma.cuh` | Split KV across warps + MMA (sm_80+) |
 | Split-Q MMA prefill | `attn_prefill_split_q_mma.cuh` | Split Q across warps + MMA (sm_80+) |
-| Paged split-KV MMA decode | `attn_paged_decode_split_kv_mma.cuh` | Paged cache + split-KV + MMA |
+
+> The paged and non-paged paths share one kernel body. Prefill is templated on
+> an independent Q schedule (`DenseQSchedule` / `PackedQSchedule`) and KV
+> source (`ContigKV` / `PagedKV`); decode only needs the KV source. There are
+> no separate `attn_paged_*.cuh` files.
 
 ### Rotary Embedding Kernel
 
@@ -26,7 +31,7 @@ The `rotary_emb` kernel (`csrc/kernels/rotary_emb.cu`) fuses cos/sin lookup and
 - One thread per (head, dim-pair), vectorized `__nv_bfloat162` load/store
 - f32 cos/sin input, bf16 compute and output
 - 256-thread blocks, grid-stride loop
-- Auto-dispatched via `apply_rotary_emb` in `astrai/extension/rotary_backend.py` (CUDA when available + inference mode, else torch complex-multiply fallback)
+- Auto-dispatched via `apply_rotary_emb` in `astrai/extension/backend/rotary.py` (CUDA when available + inference mode, else torch complex-multiply fallback)
 - No context-manager backend needed — rotary is backend-agnostic, both attention backends benefit
 
 Standalone benchmark vs torch complex-multiply (48 calls = 24 layers × q+k): 6-9x faster, max diff 0 (decode) to 3e-2 (large prefill, bf16).
@@ -50,31 +55,151 @@ CSRC_KERNELS=true pip install -e . --no-build-isolation
 # Rebuild after editing .cu/.cuh files
 CSRC_KERNELS=true python setup.py build_ext --inplace
 # Output: astrai/extension/lib/*.so
+
+# Or invoke CMake directly
+cmake -S csrc -B build/cmake \
+  -DTORCH_HOME=<site-packages>/torch \
+  -DPYTHON_INCLUDE_DIR=<python include> \
+  -DPY_SOABI=cpython-312-x86_64-linux-gnu
+cmake --build build/cmake -j 16
 ```
 
 ### Architecture flags
 
-`csrc/build.py` auto-detects the GPU compute capability and generates the appropriate `nvcc` gencode flag:
+`setup.py` passes the GPU compute capability to CMake via `ASTRAI_CUDA_ARCH` (default `89`, i.e. sm_89 / L20):
 
 - **sm_80+** (Ampere and later): enables tensor-core MMA path (`mma.sync.m16n8k16.bf16`)
 - **Below sm_80**: adds `-DASTRAI_NO_MMA` to disable the MMA path at compile time
 
 ### Build configuration
 
+`csrc/CMakeLists.txt` defines the CUDA extension build:
+
 ```
 NVCC_FLAGS = -O3 --expt-relaxed-constexpr --use_fast_math
-             --ptxas-options=-O3,-v --extra-device-vectorization --threads=8
+             --ptxas-options=-O3,-v --extra-device-vectorization --threads=16
 ```
 
-The `REGISTRY` in `csrc/build.py` lists all registered kernels (currently 4). Each entry maps a kernel name to its source files and build flags.
+Each kernel in `astrai/extension/lib` is compiled as an independent pybind11 module (one `.so` per kernel, named `<kernel>.cpython-*-x86_64-linux-gnu.so`). CMake builds all five kernel targets in parallel via `cmake --build -j N`.
+
+## Python Extension Architecture
+
+The Python extension package separates low-level kernel bindings from execution
+policy:
+
+```text
+astrai/extension/
+├── __init__.py             # Stable public API
+├── loader.py               # Optional compiled-module discovery and loading
+├── ops/
+│   ├── attention.py        # Stateless attention kernel wrappers
+│   └── rotary.py           # Stateless rotary kernel wrapper
+└── backend/
+    ├── attention.py        # Backend selection, KV cache I/O, and fallback
+    └── rotary.py           # Per-call CUDA/torch rotary dispatch
+```
+
+The dependency direction is one-way:
+
+```text
+model / inference
+       |
+       v
+extension public API
+       |
+       v
+backend policy  --->  ops wrappers  --->  loader  --->  compiled .so
+       |
+       +----------->  torch / flash-attn fallback
+```
+
+`ops` must not import `backend`. This keeps direct kernel bindings independent
+of model, cache, fallback, and backend-selection policy.
+
+### Ops Layer
+
+`astrai.extension.ops` is the low-level boundary around compiled extensions:
+
+- Wrappers are stateless and map Python arguments to pybind or
+  `torch.library.custom_op` calls.
+- Wrappers validate kernel availability and raise `RuntimeError` when a
+  requested extension was not built.
+- Wrappers do not choose another implementation, gather KV cache entries, or
+  decide whether an input is supported by a backend.
+- Tests that specifically exercise a compiled kernel may import from
+  `astrai.extension.ops`.
+
+For example, `attn_prefill(...)` means "run this CUDA kernel" rather than "run
+attention using the best available implementation":
+
+```python
+from astrai.extension.ops import attn_prefill
+
+output = attn_prefill(q, k, v, mask=mask, is_causal=True)
+```
+
+If the kernel is unavailable, this call fails. Callers that need fallback and
+capability dispatch must use the public `attention(...)` entry point instead.
+
+### Backend Layer
+
+`astrai.extension.backend` owns execution policy:
+
+- It selects CUDA, FlashAttention, or torch-native attention.
+- It checks per-call constraints such as dtype, shape, head dimension, cache
+  availability, and installed optional dependencies.
+- It owns KV cache writes and reads because those operations differ by backend.
+- It provides torch fallbacks and raises when an explicitly requested backend
+  cannot handle a call.
+- Rotary dispatch follows the same boundary without a backend class: the
+  policy layer chooses the fused op for supported inference calls and otherwise
+  uses the autograd-compatible torch implementation.
+
+Normal model and inference code should import the stable API from
+`astrai.extension`:
+
+```python
+from astrai.extension import ATTN_BACKEND, attention, attn_backend
+
+output = attention(q, k, v, kv_cache=cache, layer_id=layer_id, fwd="decode")
+
+with attn_backend(ATTN_BACKEND.TORCH_NATIVE):
+    output = attention(q, k, v)
+```
+
+The package root re-exports the supported high-level API and selected direct
+kernel wrappers. Internal code should use `astrai.extension.backend` only when
+it needs a backend type or policy implementation, and `astrai.extension.ops`
+only when it deliberately requires one exact kernel.
+
+### Placement Rules
+
+When extending this package:
+
+| Change | Location |
+|--------|----------|
+| Add a pybind call for a compiled kernel | `astrai/extension/ops/` |
+| Add argument translation required by the compiled ABI | `astrai/extension/ops/` |
+| Add capability checks or implementation selection | `astrai/extension/backend/` |
+| Add a torch or third-party fallback | `astrai/extension/backend/` |
+| Add attention KV cache behavior | `astrai/extension/backend/attention.py` |
+| Expose a supported user-facing symbol | `astrai/extension/__init__.py` |
+
+Imports belong at module scope. Optional dependencies such as `flash_attn` may
+use a module-level guarded import. Type-only imports that would create a runtime
+cycle belong under `TYPE_CHECKING`.
 
 ## Attention Backend
 
-`astrai/extension/attention_backend.py` provides the backend abstraction:
+`astrai/extension/backend/attention.py` provides the backend abstraction:
 
 - **`AttentionBackend`** (ABC): `fwd_decode` / `fwd_prefill` abstract methods, `forward` dispatches by q_len
-- **`TorchNativeBackend`**: SDPA with indirect KV cache gather (default)
-- **`CudaBackend`**: CUDA kernel dispatch — decode via `attn_paged_decode` (page_size=1), prefill via `attn_prefill`
+- **`CudaBackend`**: CUDA kernel dispatch — decode via `attn_paged_decode` (page_size=1), prefill via `attn_paged_prefill` (ragged batch, `qo_indptr` + `kv_indptr`). Default on GPU.
+- **`FlashAttnBackend`**: Optional flash-attn dispatch with `flash_attn_with_kvcache` fast path.
+- **`TorchNativeBackend`**: SDPA with indirect KV cache gather (always-available fallback)
+
+Default priority: cuda > flash > torch. Set ``ASTR_BACKEND=cuda|torch_native|flash``
+to override the default.
 
 Select a backend via context manager (mirrors `torch.nn.attention.sdpa_kernel`):
 
@@ -85,11 +210,15 @@ with attn_backend(ATTN_BACKEND.CUDA):
     engine.generate("hello")
 ```
 
-`CudaBackend` falls back to `TorchNativeBackend` when a kernel is not available.
+The `attention(...)` policy entry point falls back to `FlashAttnBackend` (when
+flash-attn is installed and supports the call) or `TorchNativeBackend` when the
+automatically selected CUDA backend cannot handle an input. An explicit
+`ASTR_BACKEND` or `attn_backend(...)` selection is strict and raises instead of
+silently switching implementations.
 
 ### Rotary Backend
 
-`astrai/extension/rotary_backend.py` provides `apply_rotary_emb(x, (cos, sin))` with auto-dispatch:
+`astrai/extension/backend/rotary.py` provides `apply_rotary_emb(x, (cos, sin))` with auto-dispatch:
 
 - **CUDA path**: calls `rotary_emb` kernel directly when available, input is bf16 on CUDA, and `torch.is_grad_enabled()` is `False` (inference)
 - **Torch fallback**: complex multiply (`torch.view_as_complex` → `torch.complex` multiply → `torch.view_as_real`), used during training (supports autograd) or when kernel unavailable
@@ -98,9 +227,9 @@ No context-manager switching needed — the dispatch is automatic per call.
 
 ## Python Wrappers
 
-`astrai/extension/attention_ops.py` provides Python wrappers for each compiled attention kernel. Each wrapper calls its CUDA kernel directly and raises `RuntimeError` if the `.so` is not available. Fallback to torch SDPA is handled by the attention backend, not the wrapper functions.
+`astrai/extension/ops/attention.py` provides Python wrappers for each compiled attention kernel. Each wrapper calls its CUDA kernel directly and raises `RuntimeError` if the `.so` is not available. Fallback to torch SDPA is handled by the attention backend, not the wrapper functions.
 
-`astrai/extension/rotary_ops.py` provides the wrapper for the rotary embedding kernel. Fallback to torch complex multiply is handled by `rotary_backend.py`.
+`astrai/extension/ops/rotary.py` provides the wrapper for the rotary embedding kernel. Fallback to torch complex multiply is handled by `backend/rotary.py`.
 
 Interface (all functions):
 ```
@@ -110,6 +239,47 @@ mask:      2D [batch, kv_len] or 3D [batch, q_len, kv_len] (bool, True=keep)
 
 Layout convention: all q/k/v are `[batch, seq_len, n_heads, head_dim]` (blhd). Scale is always `1/sqrt(head_dim)`.
 
+### Q Scheduling and KV Addressing
+
+Prefill separates Q work scheduling from KV storage:
+
+- `DenseQSchedule` maps a rectangular grid directly with
+  `batch = blockIdx.z` and `q_tile = blockIdx.x`.
+- `PackedQSchedule` consumes a compact work map for a packed
+  `[total_q, q_heads, head_dim]` tensor.
+- `ContigKV` and `PagedKV` only provide KV lengths and translate logical KV
+  positions into physical addresses. They do not schedule Q blocks.
+
+For ragged Q lengths `[70, 10, 130]` and 64 rows per Q tile, cache binding
+builds:
+
+```text
+qo_indptr       = [0, 70, 80, 210]
+q_tile_to_batch = [0, 0, 1, 2, 2, 2]
+q_tile_to_index = [0, 1, 0, 0, 1, 2]
+```
+
+Paged prefill launches:
+
+```text
+grid.x = num_q_tiles  # 6, exactly the valid ragged work items
+grid.y = q_heads
+grid.z = 1
+```
+
+Each block resolves its request and request-local tile in O(1):
+
+```cpp
+batch = q_tile_to_batch[blockIdx.x];
+q_tile = q_tile_to_index[blockIdx.x];
+```
+
+The kernel then uses `qo_indptr[batch]` for the packed Q base and adjacent
+`qo_indptr` / `kv_indptr` entries for that request's Q and KV lengths. This
+avoids the previous per-block linear scan over the batch, shared-memory
+broadcast, mapping barrier, and upper-bound grid with potentially invalid
+blocks.
+
 ## Standalone Testing
 
 Each `csrc/tests/*.cu` file has the `nvcc` compile command in its header comment. Example:
@@ -117,23 +287,22 @@ Each `csrc/tests/*.cu` file has the `nvcc` compile command in its header comment
 ```bash
 nvcc -I csrc -arch=sm_89 -O3 --use_fast_math \
      --ptxas-options=-O3,-v --extra-device-vectorization \
-     csrc/tests/attn_decode_test.cu -o /tmp/test && /tmp/test
+     -Xcompiler -fopenmp csrc/tests/attn_test.cu -o /tmp/test && /tmp/test
 ```
 
 Test files:
-- `attn_decode_test.cu` — basic decode kernel
-- `attn_paged_decode_test.cu` — paged decode kernel
-- `attn_prefill_test.cu` — prefill kernel
+- `attn_test.cu` — decode + prefill kernels (correctness tables + benchmarks)
+- `attn_paged_test.cu` — paged decode/prefill kernels
 
 ## Benchmarks
 
 Hardware: NVIDIA L20 (sm_89, 46 GB), CUDA 12.8, driver 570.86.
 
-Reproduce:
+Reproduce (decode + prefill in `attn_test.cu`, paged in `attn_paged_test.cu`):
 ```bash
 nvcc -I csrc -arch=sm_89 -O3 --use_fast_math \
      --ptxas-options=-O3,-v --extra-device-vectorization \
-     csrc/tests/attn_<name>_test.cu -o /tmp/test && /tmp/test
+     -Xcompiler -fopenmp csrc/tests/attn_test.cu -o /tmp/test && /tmp/test
 ```
 
 ## Known Optimization Targets
@@ -146,30 +315,29 @@ nvcc -I csrc -arch=sm_89 -O3 --use_fast_math \
 
 ```
 csrc/
-├── build.py              # Build system: REGISTRY, _arch_flags, nvcc flags
+├── CMakeLists.txt        # CMake build: 5 kernel targets, torch/pybind11 linking
 ├── kernels/
-│   ├── attn_common.h     # Shared attention params (AttentionParams, PagedAttentionParams)
+│   ├── attn_common.h     # Unified attention params (contig + paged modes)
 │   ├── attn_decode.cu    # Basic decode kernel (registered)
 │   ├── attn_prefill.cu   # Basic prefill kernel (registered)
 │   ├── attn_paged_decode.cu             # Paged decode kernel (registered)
+│   ├── attn_paged_prefill.cu            # Paged prefill kernel (registered)
 │   ├── rotary_emb.cu                    # Fused rotary embedding kernel (registered)
-│   ├── attn_decode_split_kv.cuh         # Split-KV variant
-│   ├── attn_decode_split_kv_mma.cuh     # Split-KV + MMA variant
-│   ├── attn_prefill_split_q.cuh         # Split-Q variant
-│   ├── attn_prefill_split_q_mma.cuh     # Split-Q + MMA variant
-│   ├── attn_paged_decode_split_kv.cuh   # Paged + split-KV variant
-│   ├── attn_paged_decode_split_kv_mma.cuh  # Paged + split-KV + MMA variant
-│   ├── attn_dispatchers.cuh             # Kernel dispatch macros
+│   ├── attn_decode_split_kv.cuh         # Split-KV variant (contig + paged via KVSource)
+│   ├── attn_decode_split_kv_mma.cuh     # Split-KV + MMA variant (contig + paged)
+│   ├── attn_prefill_split_q.cuh         # Split-Q variant (contig + paged via KVSource)
+│   ├── attn_prefill_split_q_mma.cuh     # Split-Q + MMA variant (contig + paged)
+│   ├── attn_layout_policies.cuh          # Q schedules and KVSource policies
+│   ├── attn_dispatchers.cuh             # Kernel dispatch macros + KV-templated launchers
 │   ├── attn_entry_utils.cuh             # Entry point helpers
 │   ├── attn_mma_utils.cuh               # MMA utilities
 │   └── attn_warp_utils.cuh              # Warp-level utilities
 └── tests/
     ├── test_utils.cuh           # Shared test utilities
-    ├── attn_decode_test.cu      # Decode kernel test
-    ├── attn_paged_decode_test.cu # Paged decode test
-    └── attn_prefill_test.cu     # Prefill kernel test
+    ├── attn_test.cu             # Decode + prefill kernels
+    └── attn_paged_test.cu       # Paged decode/prefill kernels
 ```
 
 Compiled `.so` files are placed in `astrai/extension/lib/`, separate from Python source files.
 
-> Document Update Time: 2026-07-31
+> Document Update Time: 2026-08-16

docs/developer/dataflow.md

@@ -14,26 +14,30 @@ This document describes the data pipeline: from raw text to model input tensors.
 ## Overview
 
 ```
-JSONL Lines → Pipeline (mask builder) → Tokenized Tensors
-                                              ↓
-                                      .h5 or .bin storage
-                                              ↓
-                                      Store.load()
+JSON / JSONL Records → Pipeline (mask builder) → Tokenized Tensors
+                                               ↓
+                                           .bin storage
+                                               ↓
+                                       Store.load()
                                               ↓
                                       Store.fetch(begin, end, keys)
                                               ↓
-                                      BaseDataset.__getitem__(idx)
-                                              ↓
-                                      Sampler → DataLoader → Training / Inference
+                                       Dataset.__getitem__(idx)
+                                               ↓
+                                       RDSampler → DataLoader → Training
 ```
 
 ## Data Preparation
 
-Raw text is tokenized via `AutoTokenizer.encode()` and saved as HDF5 (`.h5`) or binary (`.bin` + `meta.json`) files with keyed tensor groups.
+The offline `Pipeline` accepts `.jsonl` records and `.json` files containing one
+object or a list of objects. It tokenizes them and writes binary shards (`.bin`
+plus `meta.json`) with keyed tensor groups. Binary is the only registered output
+writer; the pipeline cannot emit JSONL.
 
 ### Tokenization
 
-The `Pipeline` reads JSONL lines, applies the mask builder (see [Preprocessing](../guides/preprocessing.md)), and produces flat token sequences:
+The `Pipeline` reads JSON/JSONL records, applies the mask builder (see
+[Preprocessing](../guides/preprocessing.md)), and produces token sequences:
 
 ```python
 # Per JSONL line: messages → chat template → token IDs + loss mask
@@ -42,84 +46,119 @@ loss_mask = [0, 0, 0, 1, 1, 1, 1, 1, 1]        # 0=masked, 1=train
 # Stored as flat tensors, packed with other lines by packing strategy
 ```
 
-The output `meta.json` records the storage format, key names, dtype, total token count, and tensor shapes for each shard.
+For default single-output preprocessing, the stored keys are `sequence` and
+`position_ids`, plus `loss_mask` when masking is required. Packing is supported
+for single-output data with a `sequence` key. Shard flushing counts the primary
+flat sequence for each record: `sequence` in single-output mode, otherwise the
+first flat source output.
+
+The exact shard `meta.json` schema is a top-level mapping from key to tensor
+metadata. It does not contain a storage-format or total-token field:
+
+```json
+{
+  "sequence": {"shape": [123456], "dtype": "int32"},
+  "loss_mask": {"shape": [123456], "dtype": "bool"},
+  "position_ids": {"shape": [123456], "dtype": "int32"}
+}
+```
+
+Record-aware binary data may also include `"offsets": [0, ...]` inside a key's
+metadata, but the preprocessing `BinWriter` currently does not write offsets.
 
 ### Format Detection
 
 `detect_format(load_path)` inspects the path:
 
-- If `load_path` is a file: checks suffix — `.h5`/`.hdf5` → `"h5"`, `.jsonl` → `"jsonl"`, unknown suffix raises `ValueError`
-- If `load_path` is a directory: recursively globs for `*.h5`/`*.hdf5` files → `"h5"`, `*.bin` + `**/meta.json` → `"bin"`, or `*.jsonl` + `dataset_config.json` → `"jsonl"`
+- If `load_path` is a file: `.jsonl` selects `"jsonl"`; other suffixes raise `ValueError`.
+- If `load_path` is a directory: any recursive `*.bin` plus a `meta.json` selects `"bin"`; otherwise any recursive `*.jsonl` selects `"jsonl"`.
+- Detection does not require `dataset_config.json`; configuration is selected later when `DatasetFactory.load()` constructs a transform via `_build_jsonl_transform()` and passes it to `JsonlStore.load()`.
 
 ### Store Backends
 
 Storage format is auto-detected by `detect_format()`; backends are dispatched via registry:
 
 ```
-StoreFactory.create("h5")    → H5Store
 StoreFactory.create("bin")   → MmapStore
 StoreFactory.create("jsonl") → JsonlStore
 ```
 
-All three inherit `Store` (base, owns `_data`/`_cum`/`_offsets`/`_normalize`) plus the `Streamable` and `Recordable` mixins, so every backend supports both `fetch(begin, end, keys)` (stream) and `fetch_record(index, keys)` (record) APIs.
-
-**H5Store**: Reads HDF5 files. Tensors are loaded into host memory and normalized into segmented storage. `segments_are_records=True` — each `data_i` dataset is one record.
+Both stores inherit `Store` and compose the `Streamable` and `Recordable`
+access methods.
 
 **MmapStore**: Memory-maps `.bin` files. OS page cache sharing is native — no explicit `share_memory_()` needed. Uses `torch.from_numpy(np.memmap(...))`. `segments_are_records=False` — bin segments are contiguous streams; record access is driven by `_offsets` (written when `save_bin(..., record_keys=...)` was used at preprocessing time).
 
-**JsonlStore**: On-the-fly tokenization of raw JSONL files at load time. Requires a `dataset_config.json` alongside the `.jsonl` files following the same `PipelineConfig` schema with an additional `tokenizer_path` field. Two modes: eager (default, applies `TokenizeTransform` to all records at load) and lazy (`processor=fn` given, defers tokenisation to `fetch_record` — used by DPO/GRPO).
+**JsonlStore**: Reads a `.jsonl` file or the sorted top-level `*.jsonl` files in
+a directory. Eager transform selection is owned by
+`DatasetFactory._build_jsonl_transform()` (called from `DatasetFactory.load()`) —
+the factory picks the first available route:
 
-All backends normalise tensors into `Store._data[Dict[str, List[Tensor]]]` + `Store._cum[Dict[str, List[int]]]` (cumulative lengths for bisect-based stream indexing) + `Store._offsets[Dict[str, List[int]]]` (per-record offsets for record-mode indexing). Nested keys (GRPO `responses`/`masks` as `List[List[Tensor]]`) are stored as-is and excluded from both bookkeepings — they are only accessed record-by-record.
+1. An explicit `transform=` argument passed through `store.load()`.
+2. `dataset_config.json` in the JSONL directory. It follows `PipelineConfig` and may add `tokenizer_path`; when omitted, the config directory is used.
+3. The built-in `messages` transform when `tokenizer_path=` is supplied. It masks system/user turns, trains assistant turns, and emits document-reset position IDs.
+
+`JsonlStore.load()` requires `transform=` to be passed explicitly for eager mode
+(raises `ValueError` if missing).
+
+Only DPO gets an automatic lazy route from `DatasetFactory`: raw JSONL plus
+`tokenizer_path` installs `dpo_processor` and tokenizes each record in
+`fetch_record`. GRPO does not currently have an automatic lazy processor.
+
+Eager-loaded stores normalize tensors into `Store._data[Dict[str, List[Tensor]]]` + `Store._cum[Dict[str, List[int]]]` (cumulative lengths for stream indexing) + `Store._offsets[Dict[str, List[int]]]` (per-record offsets for record indexing). Nested JSONL keys such as GRPO `responses`/`masks` are kept as record values and excluded from stream bookkeeping. Lazy DPO instead retains raw records and processes them in `fetch_record`.
 
 ## Data Keys by Training Type
 
 | Type | Storage Keys | Access Mode |
 |------|-------------|-------------|
-| `seq` | `sequence` (→ input_ids, target_ids via offset-by-1) | stream (`fetch`) |
+| `seq` | `sequence`, `position_ids` by default (`SEQDataset` consumes only `sequence`) | stream (`fetch`) |
 | `sft` | `sequence`, `loss_mask`, `position_ids` | stream (`fetch`) |
 | `dpo` | `chosen`, `rejected`, `chosen_mask`, `rejected_mask` | record (`fetch_record`) |
 | `grpo` | `prompts`, `responses`, `masks`, `rewards` | record (`fetch_record`) |
 
+Offline `.bin` output from DPO/GRPO preprocessing is not currently loadable for
+training. DPO shards are written without record offsets, while GRPO response
+groups are flattened without preserving record/group boundaries. Supported raw
+routes are eager JSONL for SEQ/SFT and automatic lazy JSONL for DPO. GRPO
+requires a caller-built, already-loaded record store.
+
 ## Dataset Architecture
 
 ```
-DatasetFactory.load(
-    train_type, load_path=None, window_size=0, stride=None,
-    storage_type=None, tokenizer_path=None,
-    max_len=2048, store=None
-)
-  → BaseDataset.load(load_path, storage_type=None)
-    → detect_format(load_path)
-    → StoreFactory.create(storage_type)
-    → Store.load(load_path)
-      → _normalize(raw)  # base Store, shared by both backends
-        → Store._data[Dict[str, List[Tensor]]]
-          + _cum[Dict[str, List[int]]]   (stream mode)
-          + _offsets[Dict[str, List[int]]]  (record mode)
+DatasetFactory.load(...)
+  → detect_format(load_path)
+  → optionally build dpo_processor for raw JSONL
+  → StoreFactory.create(storage_type, window_size, stride)
+  → Store.load(load_path, transform=... or processor=...)
+  → DatasetFactory.create(train_type, store=store)
 
 Stream datasets (SEQ/SFT):
   BaseDataset.__getitem__(idx)
-    → get_index(idx) → [begin, end)
+    → Store.sample_window(idx) → [begin, end)
     → Store.fetch(begin, end, keys) → Tensor / Dict[str, Tensor]
 
-Record datasets (DPO/GRPO via RecordDataset):
-  RecordDataset.__getitem__(idx)
+Record datasets (DPO/GRPO):
+  DPODataset/GRPODataset.__getitem__(idx)
     → Store.fetch_record(idx, keys) → Tensor / Dict[str, Tensor]
 ```
 
-Class hierarchy: `BaseDataset` ← `SEQDataset` / `SFTDataset` (stream); `BaseDataset` ← `RecordDataset` ← `DPODataset` / `GRPODataset` (record).
+Class hierarchy: `BaseDataset` is the direct base of `SEQDataset`, `SFTDataset`,
+`DPODataset`, and `GRPODataset`. There is no `RecordDataset` class.
 
 `window_size` = max input length, `stride` = step between consecutive samples (defaults to `window_size`, optional). Only meaningful for stream datasets — record datasets ignore both. `storage_type` defaults to `None` (auto-detect via `detect_format`).
 
-`tokenizer_path` triggers lazy on-the-fly tokenisation for record datasets on raw JSONL (DPO builds a `dpo_processor`; SEQ/SFT/pre-tokenised backends ignore it). `store` (pre-built `Store`) bypasses `load_path`/`storage_type`/`tokenizer_path` entirely — the caller controls Store construction.
+For raw JSONL, `tokenizer_path` builds the lazy processor only for DPO. For
+SEQ/SFT it is forwarded to `JsonlStore` so the built-in eager `messages`
+transform can be selected when no `dataset_config.json` exists. GRPO receives no
+automatic processor. A pre-built `store` bypasses path, format, tokenizer,
+window, and stride setup entirely.
 
 `Store.fetch(begin, end, keys)` (stream mode, on `Streamable`): accepts a single key (`str`) returning a `Tensor`, or a list of keys returning `Dict[str, Tensor]`. Internally uses `bisect` across multi-segment tensors. Raises `RuntimeError("Store not loaded")` if called before `load()`.
 
-`Store.fetch_record(index, keys)` (record mode, on `Recordable`): same key API. Uses `_offsets[key]` when present (bin layout with per-record offsets), otherwise indexes `_data[key]` directly (H5/JSONL where each segment is one record).
+`Store.fetch_record(index, keys)` (record mode, on `Recordable`): same key API. Uses `_offsets[key]` when present for binary record layouts; otherwise it indexes per-record JSONL tensors directly.
 
 ## Sampler
 
-`ResumableDistributedSampler` supports checkpoint-aware distributed sampling:
+`RDSampler` supports checkpoint-aware distributed sampling:
 
 - Tracks `start_epoch` / `start_iter` for resume
 - Shuffle via `torch.Generator(seed + epoch)`

docs/developer/docker-training.md

@@ -0,0 +1,57 @@
+# Containerized Training Deployment
+
+Rules for running AstrAI distributed training in containers, distilled from real deployment failures. Read before touching `Dockerfile`, `docker-compose.yml`, `scripts/train.sh`, `train-entrypoint.sh`. AGENTS.md mirrors this locally; this file is the committed version.
+
+## Architecture
+
+```
+scripts/train.sh           host-side CLI: env loading, preflight, compose wrapper, lifecycle
+  └── docker-compose.yml   GPU passthrough, mounts, in-container env vars, entrypoint
+        └── train-entrypoint.sh   GPU-count resolution, parallel-mode selection, auto-resume
+              └── train.py --config /run/astrai/train.yaml
+```
+
+| Layer | Responsible for | NOT responsible for |
+|-------|-----------------|---------------------|
+| `train.sh` | host paths, `.env.train`, preflight, lifecycle | training args, GPU selection, parallel mode |
+| compose | GPU passthrough, mounts, in-container env (NCCL) | training args (beyond `TRAIN_*` forwarding) |
+| entrypoint | `--ckpt_dir/--nprocs/--parallel_mode/--param_path`, resume | hyperparameters (YAML/CLI) |
+| `train.yaml` | hyperparameters (`_merge_yaml_into_kwargs`, CLI wins) | container paths, process count |
+
+## Path Conventions
+
+| Host var | Container | Perm | Purpose |
+|---|---|---|---|
+| `TRAIN_DATA_DIR` | `/data` | ro | dataset (`data_root_path` must be `/data`) |
+| `TRAIN_MODEL_DIR` | `/models/base` | ro | base model (`config.json` + `model.safetensors`) |
+| `TRAIN_CHECKPOINT_DIR` | `/checkpoints` | rw | checkpoint root, per-`TRAIN_JOB_NAME` subdirs |
+| `TRAIN_CONFIG_FILE` | `/run/astrai/train.yaml` | ro | training YAML (mounted only on `start`) |
+| code | `/app` | image | **not a mount**; rebuild image for code changes |
+
+## Hard Rules
+
+1. **Filter GPUs once**: compose passes the full physical set (`count: all`); `CUDA_VISIBLE_DEVICES` filters inside by physical index. Never `count: N` + physical indices (double filter leaves 1 card → `device_id out of range`).
+2. **In-container UID = host UID**: Dockerfile builds the user via `USER_UID/USER_GID` args; `train.sh` injects `ASTRAI_UID/GID` (bash `UID` is readonly). compose `user:` alone does not create the /etc/passwd entry — torch's `getpass.getuser()` then dies with `uid not found`.
+3. **In-container env vars are explicit**: `.env.train` (`--env-file`) is only compose's interpolation dictionary — never reaches the container. A var arrives only via a value-less `environment` entry (`- VAR`, read from the calling process env).
+4. **NCCL hang workaround** (this host): `NCCL_P2P_DISABLE=1` + `NCCL_NET_GDR_LEVEL=0` must be in-container.
+5. **Checkpoint complete =** `meta.json + config.json + model.safetensors + optimizer.pt + scheduler.pt`; `start` auto-resumes the latest complete one.
+6. **tqdm is silent without a TTY**: add `disable=False` in `astrai/trainer/train_callback.py`; `metric.jsonl` (per step) works as progress evidence regardless.
+
+## Operations
+
+```bash
+bash scripts/train.sh init        # first run: dirs + .env.train (edit per machine)
+bash scripts/train.sh preflight   # validate Docker/paths/GPU/model/YAML/compose
+bash scripts/train.sh start       # build + start in background (auto-resume)
+bash scripts/train.sh start --foreground -- --dry-run   # print plan only
+bash scripts/train.sh logs | status | stop | restart
+bash scripts/train.sh clean --keep 5   # prune old checkpoints (--force to delete)
+```
+
+## Files
+
+- `docker-compose.yml` — trainer service: `count: all`, `ASTRAI_UID/GID` build args + `user:`, env whitelist, mounts
+- `Dockerfile` — production stage builds user from `USER_UID/USER_GID`; `ENV HOME=/home/astrai`; `USER astrai`
+- `scripts/train.sh` — `load_env` filters `UID=` lines (readonly var); `compose()` injects `ASTRAI_UID/GID`
+- `scripts/docker/train-entrypoint.sh` — GPU-count resolution, parallel mode, resume
+- `.env.train`, `train.yaml` — host-specific; templates from `scripts/train.sh init`; scientific-notation floats (`2e-5`) parse correctly since train.py uses the YAML 1.2 float schema

docs/developer/internals.md

@@ -41,7 +41,14 @@ 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 $$
 
-`RotaryEmbedding` pre-computes `cos_table` and `sin_table` (f32, `[max_len, dim/2]`). `forward()` returns a `(cos, sin)` tuple indexed by `position_ids`. `apply_rotary_emb` applies the rotation: during training it uses torch complex multiply (autograd-compatible); during inference it auto-dispatches to a fused CUDA kernel when available. The key property is that the dot product $q_i^T k_j$ depends only on the relative position $i - j$, not the absolute positions.
+`RotaryEmbedding` pre-computes a cos/sin table `freqs_cis` of shape
+`[max_len, dim/2, 2]` (f32 — `[cos, sin]` pairs). `forward()` returns
+a `[batch, seq_len, dim/2, 2]` slice indexed by `position_ids`.
+`apply_rotary_emb` applies the rotation: during training it uses torch
+complex multiply (autograd-compatible); during inference it auto-dispatches
+to a fused CUDA kernel when available. The key property is that the dot
+product $q_i^T k_j$ depends only on the relative position $i - j$, not the
+absolute positions.
 
 **Critical for inference**: RoPE is applied **before** KV cache write, not after. If applied after caching, position encoding drift occurs because cached K/V would have stale rotation factors.
 
@@ -51,13 +58,13 @@ $$ 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_{
 
 Next-token cross-entropy with optional label smoothing:
 
-$$ L_{\text{PT}} = -\sum_{t=1}^{T} \log P(x_t \mid x_{\lt t}; \theta) $$
+$$ L_{\text{PT}} = -\frac{1}{T}\sum_{t=1}^{T} \log P(x_t \mid x_{\lt t}; \theta) $$
 
 ### SFT (Supervised Fine-Tuning)
 
 Masked cross-entropy (`ignore_index=-100`) over response tokens only:
 
-$$ L_{\text{SFT}} = -\sum_{t=P+1}^{P+L} \log P(s_t \mid s_{\lt t}; \theta) $$
+$$ L_{\text{SFT}} = -\frac{1}{L}\sum_{t=P+1}^{P+L} \log P(s_t \mid s_{\lt t}; \theta) $$
 
 Prompt tokens are masked out via `loss_mask`; only response tokens contribute to the loss.
 
@@ -81,6 +88,14 @@ Where $\rho_t = \pi_\theta(a_t|s_t) / \pi_{\text{old}}(a_t|s_t)$ is the per-toke
 
 Parameters: `group_size=4`, `clip_eps=0.2`, `kl_coef=0.01`.
 
+### MoE Load Balancing
+
+MoE layers add a differentiable load-balancing term based on mean router probabilities and top-k expert assignment frequency. The training objective is:
+
+$$ L = L_{\text{task}} + \lambda_{\text{MoE}} L_{\text{aux}} $$
+
+`TrainConfig.moe_aux_loss_coef` controls $\lambda_{\text{MoE}}$ (default `0.01`). The unweighted and weighted auxiliary losses are logged separately.
+
 ## Training Loop Internals
 
 Two-level loop: **epoch** → **batch**. Optimizer step fires every `grad_accum_steps` batches.
@@ -90,11 +105,12 @@ on_train_begin
   model.train()
   on_epoch_begin
     for batch in dataloader:
-      on_batch_begin
       with executor.accumulate(model):
-        loss = strategy.compute_loss(batch)
-        context.loss = loss.item()
-        stand_loss = loss / executor.grad_accum_steps
+        on_batch_begin
+        loss_output = strategy(batch)
+        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.consumed_samples += (
             context.config.batch_per_device * context.world_size
@@ -104,6 +120,7 @@ on_train_begin
         if executor.sync_gradients:
           on_optimizer_step
           optimizer.step()
+          strategy.on_optimizer_step()
           optimizer.zero_grad()
           if scheduler:
             scheduler.step()
@@ -112,21 +129,23 @@ on_train_end
 ```
 
 The loss is divided by `grad_accum_steps` before `backward()`, so accumulated gradients sum to the correct mean.
+Strategy metrics are detached and converted to Python `float` values before the
+`LossOutput` is returned; only `LossOutput.loss` remains a differentiable tensor.
 
 ## Callback Lifecycle
 
 | Hook | Fires | Default callback |
 |------|-------|-----------------|
-| `on_train_begin` | Before training starts | `GradientCheckpointingCallback` |
+| `on_train_begin` | Before training starts | `GradientCheckpointingCallback`, `CheckpointCallback`, `MetricCallback` |
 | `on_epoch_begin` | Start of each epoch | `ProgressBarCallback` |
 | `on_batch_begin` | Every batch | — |
-| `on_optimizer_step` | Every accumulation window | `GradientClippingCallback`, `MetricCallback`, `ProgressBarCallback` |
+| `on_optimizer_step` | Every accumulation window | `MetricCallback`, `ProgressBarCallback`, `GradientClippingCallback` |
 | `on_batch_end` | Every batch | `CheckpointCallback` |
 | `on_epoch_end` | End of each epoch | `MetricCallback`, `ProgressBarCallback` |
 | `on_error` | On exception during training | `CheckpointCallback`, `MetricCallback` |
-| `on_train_end` | Training ends (always via finally) | `CheckpointCallback`, `MetricCallback`, `GradientCheckpointingCallback` |
+| `on_train_end` | Training exits after `on_train_begin` completes (via `finally`) | `GradientCheckpointingCallback`, `CheckpointCallback`, `MetricCallback` |
 
-Default callbacks (in order): `gradient_checkpointing` (activation checkpointing, optional), `checkpoint` (safetensors, rank-0), `metric` (JSONL + validation, rank-0), `progress_bar` (tqdm), `gradient_clipping` (always registered; computes grad norm, clips only when `max_grad_norm` is not `None`).
+Default callbacks (in order): `gradient_checkpointing` (activation checkpointing, optional), `checkpoint` (safetensors, rank-0), `metric` (JSONL + validation, rank-0), `progress_bar` (tqdm, rank-0), `gradient_clipping`. The gradient-clipping callback is always registered and always calls `executor.clip_grad_norm()` with the numeric `max_grad_norm` value.
 
 ## KV Cache Mathematics
 
@@ -149,18 +168,29 @@ Three-layer separation (SGLang-inspired):
 
 - **KVStorage**: Flat token-level buffers `[n_layers, size, n_kv_heads, head_dim]`.
 - **ReqToTokenPool**: Index table `[req_idx, pos] → physical token slot`, shared across all layers.
-- **Allocator + PrefixCache**: Paged-mode slot allocation with ref-counting, LRU eviction, and hash-based prefix sharing.
+- **Allocator + RadixCache**: Paged-mode allocation with ref-counting, LRU eviction, and exact page-aligned prefix sharing when `page_size > 1`.
 
-`PagePool` orchestrates all three. In contiguous mode (default), `req_to_token` is a trivial linear mapping. In paged mode, slots are allocated on demand with prefix caching support. `bind_tasks()` returns a `KVCache` dataclass with precomputed `page_table` and `decode_mask` fields (computed once per decode step, shared across all layers). Attention layers access buffers directly — no methods, no abstraction.
+`PagePool` orchestrates all three. In contiguous mode (default), `req_to_token` is a trivial linear mapping. In paged mode, slots are allocated on demand. `RadixCache` walks exact token-page edges from the root, preserving parent-prefix context instead of treating a page hash as a globally unique key. Only complete pages whose KV entries have been materialized are shared; partial pages remain request-private and are released at completion. The final sampled token is excluded because it has not yet been decoded into KV.
+
+`bind_tasks()` returns a `KVCache` dataclass with `kv_indptr`, a prefix-sum index over sequence lengths computed once per step and shared across layers. Attention layers access buffers directly — no methods, no abstraction.
 
 ### Attention Backend
 
-Attention computation is decoupled from the model via `AttentionBackend` ABC (`astrai/extension/attention_backend.py`):
+The extension package separates mechanism from policy:
 
-- **`TorchNativeBackend`** (default): writes K/V to cache, gathers via `req_to_token` indirect indexing, calls `F.scaled_dot_product_attention`.
-- **`CudaBackend`**: decode path uses `attn_paged_decode` with `page_size=1` (the `req_to_token` table serves as the page table, each token slot is a single-token "page"); prefill path gathers K/V then calls `attn_prefill`. Falls back to `TorchNativeBackend` when kernel unavailable.
+- `astrai/extension/ops/` contains stateless wrappers that invoke one exact compiled kernel and fail when it is unavailable.
+- `astrai/extension/backend/` owns capability checks, implementation selection, fallback, and KV cache I/O.
+- Model and inference code use the stable `astrai.extension` API instead of selecting ops directly.
 
-Rotary embedding is applied via `apply_rotary_emb` in `astrai/extension/rotary_backend.py`, which auto-dispatches to the fused CUDA kernel (`rotary_emb.cu`) during inference or torch complex multiply during training (for autograd compatibility). Both attention backends share the same rotary dispatch.
+Attention computation is decoupled from the model via `AttentionBackend` ABC (`astrai/extension/backend/attention.py`):
+
+- **`CudaBackend`** (default when supported): decode path uses `attn_paged_decode` with `page_size=1` (the `req_to_token` table serves as the page table, each token slot is a single-token "page"); prefill path uses the ragged-batch `attn_paged_prefill` (addresses each request via `qo_indptr` + `kv_indptr` directly against the flat pool).
+- **`FlashAttnBackend`**: optional flash-attn dispatch with `flash_attn_with_kvcache` fast path for contiguous cache; falls back to KV gather + `flash_attn_func`.
+- **`TorchNativeBackend`** (always-available fallback): writes K/V to cache, gathers via `req_to_token` indirect indexing, calls `F.scaled_dot_product_attention`.
+- The `attention(...)` entry point uses cuda > flash > torch priority and chooses another compatible backend when an automatically selected backend cannot handle a call.
+- `ASTR_BACKEND=cuda|torch_native|flash` and `attn_backend(...)` are explicit selections; incompatible calls raise instead of silently changing backend.
+
+Rotary embedding is applied via `apply_rotary_emb` in `astrai/extension/backend/rotary.py`, which auto-dispatches to the fused CUDA kernel (`rotary_emb.cu`) during inference or torch complex multiply during training (for autograd compatibility). Both attention backends share the same rotary dispatch.
 
 Backend selection is thread-safe via `contextvars`, mirroring `torch.nn.attention.sdpa_kernel`:
 
@@ -173,6 +203,8 @@ with attn_backend(ATTN_BACKEND.CUDA):
 
 Layout convention: all q/k/v are `[batch, seq_len, n_heads, head_dim]` (blhd). Scale is always `1/sqrt(head_dim)`.
 
+Direct imports from `astrai.extension.ops` are reserved for low-level kernel tests and code that intentionally requires a specific compiled implementation. They do not provide fallback.
+
 ## Mask Algorithm Internals
 
 ### Template mode (`template: true`)
@@ -230,4 +262,4 @@ total_steps          = (batches_per_replica // grad_accum_steps) * n_epoch
 
 This accounts for data-parallel sharding — each rank processes `1/nprocs` of the dataset.
 
-> Document Update Time: 2026-07-31
+> Document Update Time: 2026-08-16

docs/get-started.md

@@ -2,11 +2,23 @@
 
 This guide walks you through installing AstrAI, downloading a model, running inference, preprocessing data, and launching your first training job.
 
+## Contents
+
+- [Prerequisites](#prerequisites)
+- [1. Install](#1-install)
+- [2. Download Model Weights](#2-download-model-weights)
+- [3. Run Inference](#3-run-inference)
+- [4. Preprocess Data](#4-preprocess-data)
+- [5. Train](#5-train)
+- [6. Evaluate](#6-evaluate)
+- [7. Docker](#7-docker)
+- [Next Steps](#next-steps)
+
 ## Prerequisites
 
 - **Python 3.12+**
-- **PyTorch 2.11+** (CUDA 12.8 recommended for GPU support)
-- NVIDIA GPU with CUDA (optional but recommended; CPU works for inference)
+- **PyTorch 2.11.0** (the exact version pinned by AstrAI; CUDA 12.8 build recommended for GPU support)
+- NVIDIA GPU with CUDA for training, `scripts/tools/generate.py`, generation evaluations, and demos. The HTTP server and direct-scoring evaluations can run on CPU where their CLI exposes a CPU device.
 
 ## 1. Install
 
@@ -24,7 +36,7 @@ pip install -e .
 # pip install -e ".[dev]"
 ```
 
-> **CUDA kernels** are opt-in. They are not built by default. When built, they can be activated via `with attn_backend(ATTN_BACKEND.CUDA):` for accelerated decode/prefill, and the fused rotary embedding kernel is auto-dispatched when available. You can skip them for normal usage.
+> **CUDA kernels** are opt-in at build time (`CSRC_KERNELS=true`). Once built, `CudaBackend` is the default attention backend on GPU (cuda > flash > torch priority). Override via `ASTR_BACKEND` env var or `attn_backend()` context manager. Fused rotary embedding kernel is auto-dispatched when available. Skip for CPU-only usage.
 
 ## 2. Download Model Weights
 
@@ -55,7 +67,7 @@ python scripts/demo/stream_chat.py
 # Type your message after >>, type !exit to quit
 ```
 
-This starts a multi-turn interactive chat session with streaming output.
+This starts a single-turn interactive prompt loop with streaming output. Each prompt is independent; conversation history is not retained.
 
 ### Start an HTTP Server
 
@@ -192,6 +204,13 @@ See [Training Guide](guides/training.md) for loss formulas and strategies. See [
 
 ## 6. Evaluate
 
+HumanEval and MMLU download their benchmark data through HuggingFace
+`datasets`, which is not part of the base install:
+
+```bash
+pip install datasets
+```
+
 ```bash
 # HumanEval (code generation, auto-downloads dataset)
 python scripts/eval/evaluate_humaneval.py --param_path ./params --num_samples 20
@@ -213,7 +232,7 @@ docker build -t astrai:latest .
 
 # Run inference server with GPU
 docker run --gpus all -p 8000:8000 astrai:latest \
-  python -m scripts.tools.server --port 8000 --device cuda
+  python scripts/tools/server.py --port 8000 --device cuda
 
 # Docker Compose (GPU)
 docker compose up -d

docs/guides/distributed.md

@@ -2,6 +2,18 @@
 
 AstrAI supports three parallel modes: **single GPU** (`none`), **Data Parallel** (`ddp`), and **Fully Sharded Data Parallel** (`fsdp`). This guide covers when to use each, how to launch multi-GPU training, and how gradient accumulation works.
 
+## Contents
+
+- [Quick Start](#quick-start)
+- [Parallel Modes](#parallel-modes)
+- [Gradient Accumulation](#gradient-accumulation)
+- [Process Launching](#process-launching)
+- [NCCL Troubleshooting](#nccl-troubleshooting)
+- [Checkpoint Saving](#checkpoint-saving)
+- [Total Steps Calculation](#total-steps-calculation)
+- [Real Examples](#real-examples)
+- [CLI Parameters](#cli-parameters)
+
 ## Quick Start
 
 ### Single GPU
@@ -21,9 +33,6 @@ python scripts/tools/train.py \
 
 ```bash
 export CUDA_VISIBLE_DEVICES=0,1,2,3
-export NCCL_P2P_DISABLE=1
-export NCCL_NET_GDR_LEVEL=0
-
 python scripts/tools/train.py \
     --train_type=sft \
     --param_path ./params \
@@ -38,9 +47,6 @@ python scripts/tools/train.py \
 
 ```bash
 export CUDA_VISIBLE_DEVICES=0,1,2,3
-export NCCL_P2P_DISABLE=1
-export NCCL_NET_GDR_LEVEL=0
-
 python scripts/tools/train.py \
     --train_type=sft \
     --param_path ./params \
@@ -110,7 +116,7 @@ AstrAI auto-detects the launch method:
 
 | Detection | Strategy | Use Case |
 |-----------|----------|----------|
-| `torchelastic` / `torchrun` env vars | `TorchrunStrategy` | External orchestrator (torchrun, SLURM, K8s) |
+| `torchelastic` / `torchrun` env vars | `TorchrunStrategy` | External orchestrator (`torchrun`, K8s) |
 | `RANK` + `WORLD_SIZE` env vars | `TorchrunStrategy` | External launch |
 | Neither | `LocalStrategy` | `python scripts/tools/train.py` (in-process spawn) |
 
@@ -126,23 +132,28 @@ For multi-node or SLURM environments:
 torchrun --nproc_per_node=4 scripts/tools/train.py \
     --train_type=sft \
     --parallel_mode=ddp \
+    --nprocs=4 \
     --param_path ./params \
     --data_root_path ./dataset \
     --batch_per_device=4
 ```
 
-When launched via torchrun, AstrAI reads `RANK`, `WORLD_SIZE`, `LOCAL_RANK` from the environment and uses `TorchrunStrategy`. The `--nprocs` flag is ignored (the orchestrator controls process count).
+When launched via `torchrun`, the launcher creates the worker processes. AstrAI reads `RANK`, `WORLD_SIZE`, and `LOCAL_RANK` from the environment and uses `TorchrunStrategy`; `--nprocs` does not control process creation in this mode.
 
-## NCCL Environment Variables
+The current training CLI still uses `--nprocs` when calculating scheduler `total_steps`. Set it to the global `WORLD_SIZE` so the step count reflects data-parallel sharding, including multi-node runs.
 
-For multi-GPU training, you **must** set these environment variables:
+Raw Slurm variables such as `SLURM_PROCID`, `SLURM_NTASKS`, and `SLURM_LOCALID` are not recognized automatically. Launch through `torchrun`, or map the scheduler's variables to `RANK`, `WORLD_SIZE`, `LOCAL_RANK`, `MASTER_ADDR`, and `MASTER_PORT` before starting AstrAI. The same requirement applies to launchers that expose only OpenMPI-specific variables.
+
+## NCCL Troubleshooting
+
+The following variables are troubleshooting options for hardware or network configurations where NCCL hangs or fails. They are not general requirements and can reduce performance by disabling peer-to-peer or GPUDirect RDMA paths:
 
 ```bash
 export NCCL_P2P_DISABLE=1
 export NCCL_NET_GDR_LEVEL=0
 ```
 
-These are required on certain hardware configurations (see `AGENTS.md`). Without them, NCCL may hang or crash during collective operations. These are set in the training shell scripts (`train-seq.sh`, `train-sft.sh`, `train-dpo.sh`) but not in Python code — you must export them before launching.
+Apply them only after confirming the relevant NCCL transport is the source of the failure. AstrAI does not set them in Python.
 
 ## Checkpoint Saving
 
@@ -176,9 +187,6 @@ This ensures the LR schedule is correctly scaled regardless of the number of GPU
 
 ```bash
 export CUDA_VISIBLE_DEVICES=0,1,2,3
-export NCCL_P2P_DISABLE=1
-export NCCL_NET_GDR_LEVEL=0
-
 python scripts/tools/train.py \
     --train_type=seq \
     --param_path ./params \
@@ -240,7 +248,7 @@ python scripts/tools/train.py \
 
 | Parameter | Default | Description |
 |-----------|---------|-------------|
-| `--nprocs` | 1 | Number of GPUs / processes |
+| `--nprocs` | 1 | Local process count for AstrAI's launcher; under `torchrun`, set it to global `WORLD_SIZE` for total-step calculation |
 | `--parallel_mode` | `fsdp` | `none`, `ddp`, or `fsdp` |
 | `--start_method` | `spawn` | Multiprocessing start method (`spawn`, `fork`, `forkserver`) |
 | `--backend` | `nccl` | Distributed backend (`nccl`, `gloo`) |
@@ -248,8 +256,8 @@ python scripts/tools/train.py \
 | `--master_port` | `29500` | Master node port |
 | `--device_type` | `cuda` | Device type |
 
-> `--tp_size` is parsed but **not yet wired** — tensor parallelism is future work. `ColumnParallelLinear` / `RowParallelLinear` exist in `astrai/parallel/module.py` but are not used by the model.
+> `--tp_size` is accepted by the CLI but discarded before configuration. Tensor parallelism is not implemented, and there is no tensor-parallel module or model integration.
 
 Full parameter reference: [CLI Reference](params.md). Training loop and strategies: [Training Guide](training.md).
 
-> Document Update Time: 2026-07-30
+> Document Update Time: 2026-08-02

docs/guides/evaluation.md

@@ -2,6 +2,29 @@
 
 AstrAI provides 7 evaluation scripts in `scripts/eval/` covering code generation, knowledge QA, perplexity, summarization, data quality, instruction following, and weight analysis.
 
+## Contents
+
+- [Prerequisites](#prerequisites)
+- [Overview](#overview)
+- [HumanEval](#humaneval-code-generation)
+- [MMLU](#mmlu-knowledge-qa)
+- [Perplexity](#perplexity-ppl)
+- [ROUGE](#rouge)
+- [IFD](#ifd-instruction-following-difficulty)
+- [IFEval](#ifeval-instruction-following)
+- [Weight Analysis](#weight-analysis)
+- [Tips](#tips)
+
+## Prerequisites
+
+HumanEval, MMLU, and IFEval import HuggingFace `datasets` to download their benchmark data. This package is not installed by AstrAI's base dependencies, so install it before running those scripts:
+
+```bash
+pip install datasets
+```
+
+The generation-based scripts require CUDA because they load the model on `cuda` with `bfloat16`. Direct-scoring and metric scripts support the devices shown below.
+
 ## Overview
 
 | Script | Metric | Model Invocation | External Dataset |
@@ -18,7 +41,15 @@ Two invocation patterns exist:
 - **Generation benchmarks** (HumanEval, IFEval): use `InferenceEngine` to generate responses, then score them.
 - **Scoring benchmarks** (MMLU, PPL, IFD): call `model()` directly under `torch.inference_mode()` for log-likelihood computation.
 
-Common defaults: `--param_path` defaults to `./params`; dtype defaults to `bfloat16` on CUDA, `float32` on CPU.
+| Script | Device support |
+|--------|----------------|
+| HumanEval | CUDA for generation; `--test_only` can score existing completions without loading a model |
+| IFEval | CUDA only |
+| MMLU | CUDA or CPU via `--device`; auto-selects CUDA when available |
+| PPL | CUDA or CPU via `--device`; auto-selects CUDA when available |
+| IFD | CUDA or CPU via `--device`; auto-selects CUDA when available |
+| ROUGE | CPU-only metric computation; no model is loaded |
+| Weight analysis | CUDA by default; CPU supported via `--device cpu` |
 
 ---
 
@@ -30,7 +61,7 @@ Generates completions for 164 programming problems, executes them against hidden
 python scripts/eval/evaluate_humaneval.py \
     --param_path ./params \
     --num_samples 20 \
-    --batch_size 32 \
+    --batch_size 64 \
     --max_tokens 512 \
     --output results/humaneval.json
 ```
@@ -47,7 +78,8 @@ python scripts/eval/evaluate_humaneval.py \
 | `--temperature` | 0.8 | Sampling temperature |
 | `--top_p` | 0.95 | Nucleus sampling threshold |
 | `--top_k` | 50 | Top-k sampling |
-| `--batch_size` | 32 | Generation batch size |
+| `--batch_size` | 64 | Generation batch size |
+| `--max_seq_len` | 4096 | KV cache sequence length |
 | `--test_workers` | 8 | ProcessPoolExecutor workers for test execution |
 | `--test_timeout` | 3.0 | Per-subprocess timeout (seconds) |
 | `--problems` | None | Restrict to specific problem indices |
@@ -66,7 +98,7 @@ python scripts/eval/evaluate_humaneval.py \
 python scripts/eval/evaluate_mmlu.py \
     --param_path ./params \
     --n_shot 5 \
-    --subjects math_algebra history_us \
+    --subjects abstract_algebra high_school_us_history \
     --output results/mmlu.json
 ```
 
@@ -82,12 +114,13 @@ python scripts/eval/evaluate_mmlu.py \
 | `--device` | auto | Device (`cuda` / `cpu`) |
 | `--dtype` | auto | `bfloat16` on CUDA, `float32` on CPU |
 | `--seed` | 0 | Seed for option permutation (0 = enabled, -1 = disabled) |
+| `--batch_size` | 4 | Questions per batch; each question produces four choice rows |
 
 **How it works**: For each question, builds a prompt with n-shot examples, then scores each choice (A/B/C/D) by computing the summed log-likelihood of the choice token given the context. The choice with the highest log-prob is the prediction.
 
 **Output**: stdout prints per-subject accuracy and overall. With `--output`, writes per-subject `{accuracy, correct, total}` + `_overall` aggregate.
 
-**Data**: Auto-downloads `cais/mmlu` from HuggingFace. Stored as per-subject CSVs in `<data_dir>/<split>/` and `<data_dir>/dev/` (for few-shot).
+**Data**: Auto-downloads `cais/mmlu` from HuggingFace. Stored as per-subject CSVs in `<data_dir>/<split>/` and `<data_dir>/dev/` (for few-shot). `--subjects` accepts canonical MMLU names such as `abstract_algebra`, `college_computer_science`, `high_school_us_history`, and `world_religions`.
 
 ---
 
@@ -100,7 +133,7 @@ python scripts/eval/evaluate_ppl.py \
     --param_path ./params \
     --input_path data.jsonl \
     --output_dir ppl_results/ \
-    --batch_size 4 \
+    --batch_size 64 \
     --max_length 2048
 ```
 
@@ -110,7 +143,7 @@ python scripts/eval/evaluate_ppl.py \
 | `--input_path` | required | Input file, glob, or directory |
 | `--output_dir` | required | Output directory for `summary.json` + token JSONL |
 | `--text_key` | `text` | Key for the text field in input data |
-| `--batch_size` | 4 | Batch size |
+| `--batch_size` | 64 | Batch size |
 | `--max_length` | 2048 | Max sequence length (tokens) |
 | `--token_level` | False | Store per-token log_probs + token-type analysis |
 | `--max_samples` | None | Random subsample per file |
@@ -119,7 +152,7 @@ python scripts/eval/evaluate_ppl.py \
 
 **Input**: JSONL or JSON files. Each item must have a field named by `--text_key` (default `text`). If `--input_path` is a directory, recursively collects `*.jsonl` and `*.json`.
 
-**Output**: `summary.json` with per-file stats (tokens, mean/median loss, perplexity, p50/p90/p95/p99). With `--token_level`, also writes per-token JSONL with token IDs and log-probs.
+**Output**: `summary.json` with per-file token count, mean loss, perplexity, and p50/p90/p95/p99 loss. Median loss is included only with `--token_level`; that mode also writes per-token JSONL with token IDs and log-probs.
 
 ---
 
@@ -181,6 +214,7 @@ python scripts/eval/evaluate_ifd.py \
 | `--sentinel_text` | `\n` | Prefix for unconditional pass (`""` → bos/pad fallback) |
 | `--per_token` | False | Include per-token IFD breakdown |
 | `--max_samples` | None | Random subsample per file |
+| `--append_eos` / `--no-append_eos` | `True` | Append (or skip) EOS token to instruction/response |
 
 **How it works**: Two forward passes per batch — (1) conditional: packed BFD sequence with context + response, (2) unconditional: response prefixed with a sentinel. IFD = mean_conditional_loss / mean_unconditional_loss. IFD > 1 means the instruction makes the response harder to predict (higher quality data).
 
@@ -210,7 +244,8 @@ python scripts/eval/evaluate_ifeval.py \
 | `--top_p` | 0.95 | Top-p sampling |
 | `--top_k` | 50 | Top-k sampling |
 | `--num_samples` | 1 | Samples per problem (best-of-n scoring) |
-| `--batch_size` | 1 | Inference batch size |
+| `--batch_size` | 64 | Inference batch size |
+| `--max_seq_len` | 4096 | KV cache sequence length |
 | `--limit` | None | Limit to first N problems (quick testing) |
 | `--dump_responses` | None | Path to dump raw responses as JSONL |
 
@@ -232,7 +267,7 @@ python scripts/eval/analyze_weights.py \
 
 | Parameter | Default | Description |
 |-----------|---------|-------------|
-| `--ckpt_dir` | required | Checkpoint dir with `model.safetensors` + `config.json` |
+| `--ckpt_dir` | required | Checkpoint directory containing `model.safetensors` |
 | `--compare` | None | Additional checkpoint dirs to compare |
 | `--no_svd` | False | Skip SVD; show only weight stats (faster) |
 | `--output` | None | Save results as JSON |
@@ -245,8 +280,8 @@ python scripts/eval/analyze_weights.py \
 ## Tips
 
 - **Quick test**: Use `--limit` (IFEval) or `--problems` (HumanEval) to run on a small subset first.
-- **Auto-download**: HumanEval, MMLU, and IFEval auto-download their datasets on first run. The other scripts expect user-provided data.
+- **Auto-download**: After installing `datasets`, HumanEval, MMLU, and IFEval auto-download their datasets on first run. The other scripts expect user-provided data.
 - **Output formats**: `--output` writes a single JSON for most scripts. PPL and IFD write an `--output_dir` containing `summary.json` plus per-file artifacts.
-- **CPU mode**: All scripts auto-detect CUDA. To force CPU, use `--device cpu --dtype float32`.
+- **CPU mode**: MMLU, PPL, and IFD support `--device cpu --dtype float32`; weight analysis supports `--device cpu`. HumanEval generation and IFEval are CUDA-only.
 
 > Document Update Time: 2026-07-30

docs/guides/inference.md

@@ -31,13 +31,17 @@ PagePool (top-level manager, orchestrates all layers)
   ├── KVStorage              k_buffer / v_buffer [n_layers, size, n_kv_heads, head_dim]
   ├── ReqToTokenPool         req_to_token [num_reqs, max_ctx_len] → physical token slot
   ├── Allocator              bitmask-based page allocator + ref-count + LRU (paged mode only)
-  └── PrefixCache            hash-based prefix matching (paged mode only)
+  └── RadixCache             exact, page-aligned prefix matching (paged mode, page_size > 1)
 ```
 
 `PagePool` supports two modes:
 
 - **Contiguous (default)**: pre-allocates `max_batch_size * max_seq_len` token slots. `req_to_token` is a trivial linear mapping (`slot = req_idx * max_seq_len + pos`). No dynamic allocation.
-- **Paged** (`page_size=1` or `>1` with `n_tokens` set): shared token pool with on-demand allocation. Allocator + PrefixCache enable prefix sharing and LRU eviction.
+- **Paged** (`page_size=1` or `>1` with `n_tokens` set): shared token pool with on-demand allocation. `Allocator` provides ref-counted allocation and LRU eviction. When `page_size > 1`, `RadixCache` also enables prefix sharing.
+
+`RadixCache` indexes complete token pages as parent-linked radix edges. Lookup walks from the root and compares each page's exact token tuple, so an identical page can only be reused under the same parent prefix. Hash values are retained for introspection, but never determine a match.
+
+Only fully materialized KV pages enter the radix. A partial final page remains private to its request and is released when the request ends. On completion, the scheduler records the prompt plus generated tokens already decoded into KV; it excludes the final sampled token because that token has not yet passed through the model. A later request resumes prefill immediately after the longest complete-page hit.
 
 `bind_tasks()` returns a `KVCache` dataclass — pure data, no methods:
 
@@ -49,22 +53,35 @@ KVCache
   ├── seq_lens               [batch_size]
   ├── out_cache_loc          [batch, seq_len] — write indices for this forward
   ├── max_len                int — max(seq_lens), avoids GPU sync in decode
-  ├── page_table             [batch, max_len] — precomputed gather indices for decode (None for prefill)
-  └── decode_mask            [batch, max_len] bool — precomputed position validity mask (None for single-batch decode)
+  ├── kv_indptr              [batch + 1] int32 — prefix sum of seq_lens, precomputed once per step
+  └── qo_indptr              [batch + 1] int32 — prefix sum of per-request q_lens (prefill), precomputed once per step
 ```
 
 Attention layers do raw buffer indexing: `k_buffer[layer_id, out_cache_loc] = k` to write, `k_buffer[layer_id, indices]` to gather.
 
 ## Attention Backend
 
+Inference code calls the policy API exported by `astrai.extension`. The
+extension implementation is split into two layers:
+
+- `astrai.extension.backend` owns capability checks, backend selection,
+  fallback, and KV cache I/O.
+- `astrai.extension.ops` contains direct wrappers around compiled CUDA kernels;
+  these wrappers raise if a kernel is unavailable and do not fall back.
+
 Attention computation (cache I/O + SDPA/kernel dispatch) is decoupled from the model via `AttentionBackend` ABC:
 
 ```
 AttentionBackend (ABC)
-  ├── TorchNativeBackend   SDPA + indirect KV cache gather (default)
-  └── CudaBackend          CUDA kernel dispatch (attn_paged_decode, attn_prefill)
+  ├── CudaBackend          CUDA kernel dispatch (default on GPU)
+  ├── FlashAttnBackend     Optional flash-attn dispatch (fallback)
+  └── TorchNativeBackend   SDPA + indirect KV cache gather (always-available fallback)
 ```
 
+Default priority is cuda > flash > torch. Automatic selection may choose a
+compatible fallback for a particular call. Set
+`ASTR_BACKEND=cuda|torch_native|flash` to require one backend process-wide.
+
 Select via context manager (mirrors `torch.nn.attention.sdpa_kernel`):
 
 ```python
@@ -74,27 +91,38 @@ with attn_backend(ATTN_BACKEND.CUDA):
     engine.generate("hello")
 ```
 
+Environment and context selections are strict: if the selected backend cannot
+handle the call, inference raises an error rather than silently switching.
+
 `CudaBackend` decode path: writes K/V to cache, then calls `attn_paged_decode` with `page_size=1` — the `req_to_token` table serves directly as the page table, each token slot is a single-token "page". No explicit K/V gather needed.
 
-`CudaBackend` prefill path: writes K/V, gathers full-sequence K/V via indirect indexing (same as `TorchNativeBackend`), then calls `attn_prefill`.
+`CudaBackend` prefill path: writes K/V, then calls `attn_paged_prefill` — a ragged-batch (paged) prefill kernel that reads K/V directly from the flat pool via `req_to_token`, addressing each request's `q_len`/`kv_len` through `qo_indptr` and `kv_indptr`. No explicit K/V gather needed.
 
-Fallback: `CudaBackend` delegates to `TorchNativeBackend` when a CUDA kernel is not available.
+Fallback: when `CudaBackend` cannot handle an input (wrong dtype or head_dim), `FlashAttnBackend` is tried next (if installed), then `TorchNativeBackend`.
+
+This fallback is performed by the public `attention(...)` policy entry point
+only when no backend was explicitly selected. Import from
+`astrai.extension.ops` only for direct kernel tests or when failure on a missing
+kernel is the intended behavior.
 
 ### Rotary Embedding Backend
 
-Rotary embedding is applied via `apply_rotary_emb` in `astrai/extension/rotary_backend.py`, which auto-dispatches:
+Rotary embedding is applied via `apply_rotary_emb` in `astrai/extension/backend/rotary.py`, which auto-dispatches:
 
-- **CUDA kernel** (`rotary_emb.cu`): fused cos/sin lookup + rotation in a single kernel, used when the kernel is available, input is on CUDA, and `torch.is_grad_enabled()` is `False` (inference mode)
+- **CUDA kernel** (`rotary_emb.cu`): fused cos/sin lookup + rotation in a single kernel, used when the kernel is available, the input is bf16 on CUDA, and `torch.is_grad_enabled()` is `False` (inference mode)
 - **Torch fallback**: complex multiply path (`torch.view_as_complex` → `torch.complex` multiply → `torch.view_as_real`), used during training (supports autograd backward) or when the CUDA kernel is not available
 
-`RotaryEmbedding` stores `cos_table`/`sin_table` as f32 buffers and returns a `(cos, sin)` tuple from `forward()`. Both attention backends share the same rotary dispatch — it is backend-agnostic.
+`RotaryEmbedding` stores a cos/sin table `freqs_cis` of shape
+`[max_len, dim/2, 2]` (f32 — `[cos, sin]` pairs) and `forward()` returns
+a `[batch, seq_len, dim/2, 2]` slice indexed by `position_ids`. Both
+attention backends share the same rotary dispatch — it is backend-agnostic.
 
 ## Continuous Batching
 
 `InferenceScheduler` runs a daemon thread with a 4-phase loop:
 
 ```
-1. Cleanup → Remove finished tasks, free KV cache slots/pages
+1. Cleanup → Record complete materialized pages, then release task-owned KV resources
 2. Refill  → Pop from waiting_queue, task_alloc resources, activate
 3. Prefill → Group by (prompt_len, start_pos), run full forward
 4. Decode  → Run single-token forward for each same-position group
@@ -136,7 +164,6 @@ Adding a protocol = one builder file, no handler subclassing needed.
 ```
 InferenceEngine
   ├── generate(prompt, stream, ...) → str | List[str] | Generator
-  ├── generate_with_request(req)    → same
   ├── generate_async(prompt, ...)   → AsyncGenerator
   ├── get_stats()                   → Dict
   └── shutdown()
@@ -183,22 +210,45 @@ curl -X POST http://localhost:8000/v1/messages \
   -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`.
+Supports `stop_sequences` and streaming via `event: content_block_delta`. Anthropic streams also end with the shared `data: [DONE]` sentinel after `event: message_stop`.
 
-### GenerationRequest Parameters
+### Request Parameters
+
+The HTTP protocols and direct engine API have distinct request models and defaults.
+
+**OpenAI** (`ChatCompletionRequest`):
 
 | Param | Type | Default | Description |
 |-------|------|---------|-------------|
+| `model` | str | `"astrai"` | Model name returned in responses |
 | `messages` | List[dict] | required | Chat messages (role, content) |
-| `top_k` | int | 50 | Top-k count |
-| `top_p` | float | 1.0 | Nucleus threshold |
-| `temperature` | float | 1.0 | Sampling temperature (> 0.0) |
-| `max_tokens` | Optional[int] | None | Max generation length |
-| `stream` | bool | False | Stream output |
+| `temperature` | Optional[float] | 1.0 | Sampling temperature (0.0-2.0) |
+| `top_p` | Optional[float] | 1.0 | Nucleus threshold (0.0-1.0) |
+| `top_k` | Optional[int] | 50 | Top-k count |
+| `max_tokens` | Optional[int] | 2048 | Max generation length |
+| `stream` | Optional[bool] | False | Stream output |
 | `stop` | Optional[Union[str, List[str]]] | None | Stop sequences |
-| `frequency_penalty` | float | 0.0 | Frequency penalty |
-| `tools` | Optional[List[dict]] | None | Tool definitions for function calling |
-| `tool_choice` | Optional[str] | None | Tool selection mode |
+| `n` | Optional[int] | 1 | Number of choices requested |
+| `presence_penalty` | Optional[float] | 0.0 | Presence penalty (-2.0 to 2.0) |
+| `frequency_penalty` | Optional[float] | 0.0 | Frequency penalty (-2.0 to 2.0) |
+| `logit_bias` | Optional[Dict[int, float]] | None | Per-token logit bias |
+| `user` | Optional[str] | None | End-user identifier |
+| `tools` | Optional[List[ToolDef]] | None | Tool definitions for function calling |
+| `tool_choice` | Optional[Union[str, Dict[str, Any]]] | `"auto"` | Tool selection mode or explicit tool choice |
+
+**Anthropic** (`MessagesRequest`):
+
+| Param | Type | Default | Description |
+|-------|------|---------|-------------|
+| `model` | str | `"astrai"` | Model name returned in responses |
+| `messages` | List[AnthropicMessage] | required | User/assistant messages |
+| `system` | Optional[str] | None | System prompt |
+| `max_tokens` | int | 1024 | Max generation length |
+| `temperature` | Optional[float] | 1.0 | Sampling temperature (0.0-2.0) |
+| `top_p` | Optional[float] | 1.0 | Nucleus threshold (0.0-1.0) |
+| `top_k` | Optional[int] | 50 | Top-k count |
+| `stream` | Optional[bool] | False | Stream output |
+| `stop_sequences` | Optional[List[str]] | None | Stop sequences |
 
 ### SSE Streaming Format
 
@@ -240,6 +290,8 @@ data: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":
 
 event: message_stop
 data: {"type":"message_stop"}
+
+data: [DONE]
 ```
 
 ### Error Responses
@@ -294,4 +346,4 @@ async for token in engine.generate_async("Hello", ...):    # -> AsyncGenerator[s
     print(token)
 ```
 
-> Document Update Time: 2026-07-31
+> Document Update Time: 2026-08-16