wehub/nvlabs--longlive
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:31:40 +08:00
commit 4df75a30cd
234 changed files with 41407 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+41
View File
@@ -0,0 +1,41 @@
.DS_Store
__pycache__/
*.py[cod]
*$py.class
.pytest_cache/
.mypy_cache/
.ruff_cache/
longlive_models/
wan_models/
videos/
logs/
outputs/
wandb/
debug_save/
pretrained/
.codex/
.claude/
# Local reports and generated release artifacts
*.html
*.log
*.out
*.err
*.mp4
*.pt
*.pth
# Optimization experiment records
agent/logs/
agent/runs/
agent/profiles/
# Native and build artifacts
*.so
*.egg-info/
build/
dist/
fouroversix/build/
utils/kernel/build/
CONTRIBUTING.md
+141
View File
@@ -0,0 +1,141 @@
## LongLive OSS Contribution Rules
#### Issue Tracking
* All enhancement, bugfix, or change requests must begin with the creation of a [LongLive Issue Request](https://github.com/nvidia/LongLive/issues).
* The issue request must be reviewed by LongLive engineers and approved prior to code review.
#### Coding Guidelines
- All source code contributions must strictly adhere to the [LongLive Coding Guidelines](CODING-GUIDELINES.md).
- In addition, please follow the existing conventions in the relevant file, submodule, module, and project when you add new code or when you extend/fix existing functionality.
- To maintain consistency in code formatting and style, you should also run `clang-format` on the modified sources with the provided configuration file. This applies LongLive code formatting rules to:
- class, function/method, and variable/field naming
- comment style
- indentation
- line length
- Format git changes:
```bash
# Commit ID is optional - if unspecified, run format on staged changes.
git-clang-format --style file [commit ID/reference]
```
- Format individual source files:
```bash
# -style=file : Obtain the formatting rules from .clang-format
# -i : In-place modification of the processed file
clang-format -style=file -i -fallback-style=none <file(s) to process>
```
- Format entire codebase (for project maintainers only):
```bash
find samples plugin -iname *.h -o -iname *.c -o -iname *.cpp -o -iname *.hpp \
| xargs clang-format -style=file -i -fallback-style=none
```
- Avoid introducing unnecessary complexity into existing code so that maintainability and readability are preserved.
- Try to keep pull requests (PRs) as concise as possible:
- Avoid committing commented-out code.
- Wherever possible, each PR should address a single concern. If there are several otherwise-unrelated things that should be fixed to reach a desired endpoint, our recommendation is to open several PRs and indicate the dependencies in the description. The more complex the changes are in a single PR, the more time it will take to review those changes.
- Write commit titles using imperative mood and [these rules](https://chris.beams.io/posts/git-commit/), and reference the Issue number corresponding to the PR. Following is the recommended format for commit texts:
```
#<Issue Number> - <Commit Title>
<Commit Body>
```
- Ensure that the build log is clean, meaning no warnings or errors should be present.
- Ensure that all `sample_*` tests pass prior to submitting your code.
- All OSS components must contain accompanying documentation (READMEs) describing the functionality, dependencies, and known issues.
- See `README.md` for existing samples and plugins for reference.
- All OSS components must have an accompanying test.
- If introducing a new component, such as a plugin, provide a test sample to verify the functionality.
- To add or disable functionality:
- Add a CMake option with a default value that matches the existing behavior.
- Where entire files can be included/excluded based on the value of this option, selectively include/exclude the relevant files from compilation by modifying `CMakeLists.txt` rather than using `#if` guards around the entire body of each file.
- Where the functionality involves minor changes to existing files, use `#if` guards.
- Make sure that you can contribute your work to open source (no license and/or patent conflict is introduced by your code). You will need to [`sign`](#signing-your-work) your commit.
- Thanks in advance for your patience as we review your contributions; we do appreciate them!
#### Pull Requests
Developer workflow for code contributions is as follows:
1. Developers must first [fork](https://help.github.com/en/articles/fork-a-repo) the [upstream](https://github.com/nvidia/LongLive) LongLive OSS repository.
2. Git clone the forked repository and push changes to the personal fork.
```bash
git clone https://github.com/YOUR_USERNAME/YOUR_FORK.git LongLive
# Checkout the targeted branch and commit changes
# Push the commits to a branch on the fork (remote).
git push -u origin <local-branch>:<remote-branch>
```
3. Once the code changes are staged on the fork and ready for review, a [Pull Request](https://help.github.com/en/articles/about-pull-requests) (PR) can be [requested](https://help.github.com/en/articles/creating-a-pull-request) to merge the changes from a branch of the fork into a selected branch of upstream.
* Exercise caution when selecting the source and target branches for the PR.
Note that versioned releases of LongLive OSS are posted to `release/` branches of the upstream repo.
* Creation of a PR creation kicks off the code review process.
* Atleast one LongLive engineer will be assigned for the review.
* While under review, mark your PRs as work-in-progress by prefixing the PR title with [WIP].
4. Since there is no CI/CD process in place yet, the PR will be accepted and the corresponding issue closed only after adequate testing has been completed, manually, by the developer and/or LongLive engineer reviewing the code.
#### Signing Your Work
* We require that all contributors "sign-off" on their commits. This certifies that the contribution is your original work, or you have rights to submit it under the same license, or a compatible license.
* Any contribution which contains commits that are not Signed-Off will not be accepted.
* To sign off on a commit you simply use the `--signoff` (or `-s`) option when committing your changes:
```bash
$ git commit -s -m "Add cool feature."
```
This will append the following to your commit message:
```
Signed-off-by: Your Name <your@email.com>
```
* Full text of the DCO:
```
Developer Certificate of Origin
Version 1.1
Copyright (C) 2004, 2006 The Linux Foundation and its contributors.
1 Letterman Drive
Suite D4700
San Francisco, CA, 94129
Everyone is permitted to copy and distribute verbatim copies of this license document, but changing it is not allowed.
```
```
Developer's Certificate of Origin 1.1
By making a contribution to this project, I certify that:
(a) The contribution was created in whole or in part by me and I have the right to submit it under the open source license indicated in the file; or
(b) The contribution is based upon previous work that, to the best of my knowledge, is covered under an appropriate open source license and I have the right under that license to submit that work with modifications, whether created in whole or in part by me, under the same open source license (unless I am permitted to submit under a different license), as indicated in the file; or
(c) The contribution was provided directly to me by some other person who certified (a), (b) or (c) and I have not modified it.
(d) I understand and agree that this project and the contribution are public and that a record of the contribution (including all personal information I submit with it, including my sign-off) is maintained indefinitely and may be redistributed consistent with this project or the open source license(s) involved.
```
LICENSE
+201
View File
@@ -0,0 +1,201 @@
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
+282
View File
@@ -0,0 +1,282 @@
<p align="center" style="border-radius: 10px">
<img src="assets/longlive2/logo.png" width="100%" alt="LongLive2.0 logo"/>
</p>
# 🎬 LongLive 2.0: An NVFP4 Parallel Infrastructure for Long Video Generation
[![Paper](https://img.shields.io/badge/Paper-LongLive_2.0-brown)](https://arxiv.org/abs/2605.18739)
[![Paper](https://img.shields.io/badge/Paper-LongLive_1.0-orange)](https://github.com/NVlabs/LongLive/tree/v1.0)
[![Paper](https://img.shields.io/badge/Paper-LongLive_RAG-yellow)](https://github.com/qixinhu11/LongLive-RAG)
[![Video](https://img.shields.io/badge/YouTube-Video-red)](https://www.youtube.com/watch?v=7oQALy32fiU)
[![Code](https://img.shields.io/badge/GitHub-Code-blue)](https://github.com/NVlabs/LongLive)
[![Demo](https://img.shields.io/badge/Demo-Page-green)](https://nvlabs.github.io/LongLive/LongLive2/)
[![Docs](https://img.shields.io/badge/Full-Documentation-brightgreen)](https://nvlabs.github.io/LongLive/LongLive2/docs/)
<div align="center">
<!-- TODO: replace this text block with the final project-page video/demo embed. -->
[![Watch the video](assets/longlive2/first-video-frame.png)](https://www.youtube.com/watch?v=7oQALy32fiU)
</div>
## 💡 TLDR: Infra with NVFP4 and parallelism for both training and inference
<p align="center" style="border-radius: 10px">
<img src="assets/longlive2/teaser.jpg" width="100%" alt="LongLive2.0 teaser"/>
</p>
## News
- 🔥 [2026.07.08] We support FP8 inference on LongLive 2.0. Please refer to [here](https://github.com/NVlabs/LongLive#fp8-ptq).
- 🔥 [2026.06.01] We released [LongLive-RAG](https://github.com/qixinhu11/LongLive-RAG), a general retrieval-augmented framework for long video gen.
- 🔥 [2026.05.30] LongLive2.0 now supports I2V AR teacher-forcing training and I2V DMD distillation for Wan2.2-TI2V-5B.
- ⚡ [2026.05.25] We optimized the NVFP4 inference path with fused Triton RoPE/adaLN kernels, reduced KV-cache synchronization overhead, in-place quantized KV-cache updates, faster FP4 KV dequantization, pinned VAE transfers, and safer LoRA-before-quantization setup, improving overall throughput by **18.6%**.
- 🔥 [2026.05.13] We release **LongLive 2.0**, infra with NVFP4, parallelism and multi-shot for AR training, DMD distillation, and inference (⚡45.7 FPS). The original LongLive 1.0 is now in the [v1.0](https://github.com/NVlabs/LongLive/tree/v1.0) branch.
- 🔥 [2026.04.12] LongLive supports kv cache compression with [TriAttention](https://github.com/WeianMao/triattention), with 50% KV reduction and no quality drop. Check it [here](https://github.com/WeianMao/triattention/tree/main/longlive)
- 🎉 [2026.1.27] LongLive is accepted by **ICLR-2026**.
- 🔥 [2026.1.11] LongLive supports adapting LongLive's original RoPE into KV-cache relative RoPE and generates infinite long videos!
- 🔥 [2025.11.3] We implement LongLive on linear attention model [SANA-Video](https://nvlabs.github.io/Sana/Video/)! Now SANA-Video can generate 60s interactive videos in real-time.
- 🔥 [2025.9.29] We release [Paper](https://arxiv.org/abs/2509.22622), this GitHub repo [LongLive](https://github.com/NVlabs/LongLive) with all training and inference code, the model weight [LongLive-1.3B](https://huggingface.co/Efficient-Large-Model/LongLive-1.3B), and demo page [Website](https://nvlabs.github.io/LongLive).
## Introduction
**LongLive 1.0**: Real-time Interactive Long Video Generation. [You can find it here](https://github.com/NVlabs/LongLive/tree/v1.0) in our V1.0 branch.
**LongLive 2.0**: an NVFP4 Parallel Infrastructure for Long Video Generation
- For training, it supports
- [x] Balanced sequence parallel for T2V/I2V AR training (teacher-forcing).
- [x] T2V/I2V AR training on multi-shot (or single-shot) videos.
- [x] NVFP4 (or BF16) for both AR training and few-step distillation.
- For inference, it supports
- [x] NVFP4 inference (W4A4) and NVFP4 KV Cache.
- [x] TorchAO FP8 PTQ inference (W8A8) from the BF16 checkpoint.
- [x] Multi-shot attention sink.
- [x] Sequence parallel inference.
- [x] Async decoding.
<p align="left" style="border-radius: 10px">
<img src="assets/longlive2/fig_framework_overview.png" width="80%" alt="LongLive2.0 framework overview"/>
</p>
**LongLive 1.0**: Real-time Interactive Long Video Generation. It accepts sequential user prompts and generates corresponding videos in real time, enabling user-guided long video generation. The key insights are attention sink, KV-recache, and streaming long tuning.
<p align="left" style="border-radius: 10px">
<img src="assets/longlive2/LongLive1_teaser.png" width="80%" alt="LongLive1.0 framework overview"/>
</p>
## Getting Started
- [Full Documentation](https://nvlabs.github.io/LongLive/LongLive2/docs/)
- [Installation](https://nvlabs.github.io/LongLive/LongLive2/docs/#installation)
- [NVFP4 Setup](https://nvlabs.github.io/LongLive/LongLive2/docs/#nvfp4-installation)
- [Training Modes](https://nvlabs.github.io/LongLive/LongLive2/docs/#training)
- [Inference](https://nvlabs.github.io/LongLive/LongLive2/docs/#inference)
- [Data Organization](https://nvlabs.github.io/LongLive/LongLive2/docs/#training-data)
The default git clone fetches objects from all branches, including our demopage branch, which contains large assets. For normal use, only the main branch is needed. Please clone only main with:
```git clone --single-branch --branch main --depth 1 https://github.com/NVlabs/LongLive.git```
### Quick Start
#### BF16
```python
import torch
from omegaconf import OmegaConf
from pipeline import CausalDiffusionInferencePipeline
from utils.config import normalize_config
from utils.inference_utils import (
load_generator_checkpoint,
place_vae_for_streaming,
prepare_single_prompt_inputs,
save_video,
)
prompt = "A compact silver robot walks through a clean robotics lab."
merged_checkpoint_path = "LongLive-2.0-5B/model_bf16.pt"
config = normalize_config(OmegaConf.load("configs/inference.yaml"))
device = torch.device("cuda")
torch.set_grad_enabled(False)
pipe = CausalDiffusionInferencePipeline(config, device=device)
load_generator_checkpoint(pipe.generator, merged_checkpoint_path)
pipe = pipe.to(device=device, dtype=torch.bfloat16)
place_vae_for_streaming(pipe, config) # honor streaming_vae + vae_device when set
pipe.generator.model.eval().requires_grad_(False)
noise, prompts = prepare_single_prompt_inputs(config, prompt, device)
video = pipe.inference(noise=noise, text_prompts=prompts)
save_video(video[0], "videos/quickstart/sample.mp4", fps=24)
```
`place_vae_for_streaming` is a no-op unless `inference.streaming_vae` is true and `inference.vae_device` is set, so toggling streaming-pipeline decode in your yaml is enough — the script does not need to change.
#### FP8 PTQ
Download `model_bf16.pt` from
[`Efficient-Large-Model/LongLive-2.0-5B`](https://huggingface.co/Efficient-Large-Model/LongLive-2.0-5B),
set `checkpoints.generator_ckpt` in `configs/fp8/inference_fp8.yaml`, and run:
```bash
python inference.py --config_path configs/fp8/inference_fp8.yaml
```
This loads the BF16 generator, applies TorchAO row-wise dynamic FP8 W8A8 PTQ,
and then enables the existing `torch.compile` path. With the provided 5B model,
300 eligible core Linear layers use FP8; six small conditioning/output
projections stay in BF16 for stability and to avoid FP8 overhead.
The validated stack is Python 3.10, PyTorch 2.8.0+cu128, and TorchAO 0.13.0 on
H100 (SM90); compute capability 8.9 or newer is required. The supplied config
uses `torch_compile: auto`. Its `max-autotune` warm-up can take several minutes
while guard/shape variants are compiled, so use repeated inference and exclude
all compile/warm-up samples when measuring steady-state performance. Set
`torch_compile: false` for a short eager-mode smoke test.
The supplied config uses the single 8-latent-frame block validated on H100.
Longer generation introduces additional KV-cache shapes and may trigger more
compilation or eager fallback; validate the intended frame count before
benchmarking or deployment.
#### NVFP4
Point `checkpoints.generator_ckpt` in `configs/nvfp4/inference_nvfp4.yaml` at the downloaded checkpoint and set `model_quant_use_transformer_engine` according to the backend you are using:
- TransformerEngine checkpoint (`model_te.pt`): `model_quant_use_transformer_engine: true`
- FourOverSix checkpoint (`model_4o6.pt`): `model_quant_use_transformer_engine: false`
`setup_nvfp4_pipeline` handles checkpoint loading, NVFP4 module wrapping, weight materialization, dtype/device placement, and the streaming-pipeline VAE relocation for both backends — the bf16 `pipe.to(...)` shortcut is unsafe here because it would cast the quantized buffers.
```python
import torch
from omegaconf import OmegaConf
from pipeline import CausalDiffusionInferencePipeline
from utils.config import normalize_config
from utils.inference_utils import prepare_single_prompt_inputs, save_video, setup_nvfp4_pipeline
prompt = "A compact silver robot walks through a clean robotics lab."
config = normalize_config(OmegaConf.load("configs/nvfp4/inference_nvfp4.yaml"))
device = torch.device("cuda")
torch.set_grad_enabled(False)
pipe = CausalDiffusionInferencePipeline(config, device=device)
setup_nvfp4_pipeline(pipe, config, device)
pipe.generator.model.eval().requires_grad_(False)
noise, prompts = prepare_single_prompt_inputs(config, prompt, device)
video = pipe.inference(noise=noise, text_prompts=prompts)
save_video(video[0], "videos/quickstart/sample_nvfp4.mp4", fps=24)
```
## Training Modes
LongLive2.0 supports both T2V and I2V training. Each modality follows the same two-stage recipe: AR teacher-forcing training first, then DMD distillation from the AR checkpoint.
### T2V Training
```bash
torchrun --standalone --nnodes=1 --nproc_per_node=8 train.py \
--config_path configs/train_ar.yaml \
--logdir logs/train_ar \
--wandb-save-dir wandb \
--disable-wandb
torchrun --standalone --nnodes=1 --nproc_per_node=8 train.py \
--config_path configs/train_dmd.yaml \
--logdir logs/train_dmd \
--wandb-save-dir wandb \
--disable-wandb
```
### I2V Training
```bash
torchrun --standalone --nnodes=1 --nproc_per_node=8 train.py \
--config_path configs/train_i2v_ar.yaml \
--logdir logs/train_i2v_ar \
--wandb-save-dir wandb \
--disable-wandb
torchrun --standalone --nnodes=1 --nproc_per_node=8 train.py \
--config_path configs/train_i2v_dmd.yaml \
--logdir logs/train_i2v_dmd \
--wandb-save-dir wandb \
--disable-wandb
```
For I2V configs, set `algorithm.i2v: true` and `algorithm.independent_first_frame: true`. `data.image_or_video_shape[1]` is the full latent sequence length, for example `96`, not `96 + 1`: the clean image latent replaces the first latent during denoising and that first latent is masked out of the training loss. For I2V DMD, set `checkpoints.generator_ckpt` to the I2V AR checkpoint used to initialize the student.
## Models
| Model | FPS ↑ | Params | VBench ↑ | Multi-shot |
| --- | ---: | ---: | ---: | :---: |
| [LongLive-1.3B](https://huggingface.co/Efficient-Large-Model/LongLive-1.3B) | 20.7 | 1.3B | 84.87 | |
| [LongLive-2.0-5B](https://huggingface.co/Efficient-Large-Model/LongLive-2.0-5B) | 24.8 | 5B | 85.06 | ✅ |
| [LongLive-2.0-5B-NVFP4-4Step](https://huggingface.co/Efficient-Large-Model/LongLive-2.0-5B-NVFP4-S4) | 29.7 | 5B | 84.51 | ✅ |
| [LongLive-2.0-5B-NVFP4-2Step](https://huggingface.co/Efficient-Large-Model/LongLive-2.0-5B-NVFP4-S2) | 45.7 | 5B | 83.14 | ✅ |
## Awesome work using LongLive
- [DreamForge-World 0.1](https://trydreamforge.com/): Adapts the LongLive AR video stack with a residual action pathway for low-compute real-time controllable world modeling.
- [DreamX-World 1.0](https://arxiv.org/abs/2606.16993): Follows LongLive by adapting the model on long sequences with long rollouts and local temporal windows for stable long-horizon AR world generation.
- [SANA-Video](https://nvlabs.github.io/Sana/docs/longsana/): Combines SANA-Video with LongLive to build LongSANA, a real-time minute-long video generation variant with constant-memory KV cache.
- [Daydream Scope](https://docs.daydream.live/scope/reference/pipelines/longlive): Wraps LongLive as a streaming AR video diffusion pipeline for interactive text-to-video and video-to-video workflows.
- [MemFlow](https://github.com/KlingAIResearch/MemFlow): Builds on the LongLive codebase and adds adaptive memory retrieval for more consistent long narrative video generation.
- [ShotStream](https://github.com/KlingAIResearch/ShotStream): Builds on LongLive’s distillation procedure for real-time streaming multi-shot AR video generation.
- [Stream-T1](https://github.com/FrameX-AI/Stream-T1): Builds on LongLive's codebase and algorithm, adding test-time scaling with noise propagation, reward pruning, and memory sinking.
- [KVPO](https://github.com/Richard-Zhang-AI/KVPO): Builds on LongLive and related AR video codebases to perform GRPO-style alignment through historical KV semantic exploration.
- [LoL](https://github.com/justincui03/LoL): Builds on LongLive to study and mitigate sink-collapse for ultra-long AR streaming video generation.
- [TriAttention](https://github.com/WeianMao/triattention/tree/main/longlive): Integrates trigonometric KV-cache compression into LongLive's causal inference pipeline, reducing KV memory inside LongLive's local-attention window.
- [StreamEdit](https://github.com/DSL-Lab/StreamEdit): Provides a `LongLive_StreamEdit` implementation for training-free streaming video editing built on the LongLive v1.0 codebase.
- [Streaming Autoregressive Video Generation via Diagonal Distillation](https://github.com/Sphere-AI-Lab/diagdistill): Builds on the LongLive codebase and supports direct initialization from `LongLive-1.3B` checkpoints for streaming AR video distillation.
- [Forcing-KV](https://github.com/zju-jiyicheng/Forcing-KV): Adds hybrid KV-cache compression to LongLive, including LongLive inference and interactive-generation scripts.
- [Dummy Forcing](https://github.com/csguoh/DummyForcing): Unifies Self-Forcing, LongLive, and Causal-Forcing pipelines with LongLive inference, VBench, and interactive-generation configs.
- [MemRoPE](https://github.com/YoungRaeKimm/MemRoPE): Uses LongLive as a supported base model for training-free infinite video generation with evolving memory tokens.
- [Astrolabe](https://github.com/franklinz233/Astrolabe): Supports LongLive as a distilled autoregressive video backbone with LongLive-specific RL configs and LoRA initialization.
## License
This repository is released under the Apache 2.0 license. See [LICENSE](LICENSE) for details.
## Citation
Please consider citing our work if you find them useful:
```bibtex
@article{longlive_2.0,
title={LongLive2.0: An NVFP4 Parallel Infrastructure for Long Video Generation},
author={Chen, Yukang and Wang, Luozhou and Huang, Wei and Yang, Shuai and Zhang, Bohan and Xiao, Yicheng and Chu, Ruihang and Mao, Weian and Hu, Qixin and Liu, Shaoteng and Zhao, Yuyang and Mao, Huizi and Chen, Ying-Cong and Xie, Enze and Qi, Xiaojuan and Han, Song},
journal={arXiv preprint arXiv: 2605.18739},
year={2026}
}
```
```bibtex
@inproceedings{longlive,
title={Longlive: Real-time interactive long video generation},
author={Yang, Shuai and Huang, Wei and Chu, Ruihang and Xiao, Yicheng and Zhao, Yuyang and Wang, Xianbang and Li, Muyang and Xie, Enze and Chen, Yingcong and Lu, Yao and others},
booktitle={ICLR},
year={2026},
}
```
```bibtex
@article{longlive_rag,
title = {LongLive-RAG: A General Retrieval-Augmented Framework for Long Video Generation},
author = {Hu, Qixin and Yang, Shuai and Huang, Wei and Han, Song and Chen, Yukang},
journal = {arXiv preprint arXiv:2606.02553},
year = {2026}
}
```
## Acknowledgement
- [Self-Forcing](https://github.com/guandeh17/Self-Forcing): the AR training codebase and formulation we build upon.
- [Wan2.2](https://github.com/Wan-Video/Wan2.2): the base video diffusion model components used in this release.
README.wehub.md
+7
View File
@@ -0,0 +1,7 @@
# WeHub 来源说明
- 原始项目：`NVlabs/LongLive`
- 原始仓库：https://github.com/NVlabs/LongLive
- 导入方式：上游默认分支的最新快照
- 原作者、版权和许可证信息以原始仓库及本仓库 LICENSE 为准
- 本文件仅用于记录来源，不代表 WeHub 是原项目作者
assets/longlive2/LongLive1_teaser.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 1.1 MiB

assets/longlive2/fig_framework_overview.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 115 KiB

assets/longlive2/first-video-frame.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 1.8 MiB

assets/longlive2/logo.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 183 KiB

assets/longlive2/teaser.jpg
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 519 KiB

configs/fp8/inference_fp8.yaml
+51
View File
@@ -0,0 +1,51 @@
# TorchAO FP8 PTQ inference from the merged BF16 LongLive checkpoint.
model_kwargs:
model_name: Wan2.2-TI2V-5B
timestep_shift: 5.0
num_frame_per_block: 8
local_attn_size: 32
use_ema: false
output_folder: videos/longlive2_fp8
num_samples: 1
num_output_frames: 8
save_latents_only: false
inference_iter: -1
data:
data_path: example/long_example.txt
image_or_video_shape:
- 1
- 8
- 48
- 44
- 80
inference:
sampling_steps: 4
sink_size: 8
guidance_scale: 1.0
multi_shot_sink: true
multi_shot_rope_offset: 8
streaming_vae: false
async_vae: false
vae_type: wan
checkpoints:
generator_ckpt: /path/to/LongLive-2.0-5B/model_bf16.pt
# Load the checkpoint in BF16, then quantize compatible generator Linear layers
# to row-wise dynamic FP8 (W8A8) before torch.compile.
fp8_quant: true
torch_compile: auto
torch_compile_min_samples: 3
torch_compile_target: generator_model
torch_compile_backend: inductor
torch_compile_mode: max-autotune-no-cudagraphs
torch_compile_fullgraph: false
torch_compile_dynamic: false
torch_compile_suppress_errors: false
logging:
seed: 0
configs/inference.yaml
+48
View File
@@ -0,0 +1,48 @@
# Passed directly to WanDiffusionWrapper(**model_kwargs).
model_kwargs:
model_name: Wan2.2-TI2V-5B
timestep_shift: 5.0
num_frame_per_block: 8
local_attn_size: 32
use_ema: false
output_folder: videos/longlive2
num_samples: 1
save_latents_only: false
data:
data_path: /path/to/longlive2/inference_prompts
# Latent tensor shape [B, F, C, H, W]. F is also the default number of
# generated latent frames unless num_output_frames is provided.
image_or_video_shape:
- 1
- 128
- 48
- 44
- 80
inference:
sampling_steps: 4
sink_size: 8
guidance_scale: 1.0
multi_shot_sink: true
multi_shot_rope_offset: 8
streaming_vae: false
async_vae: false
vae_type: wan # wan, mg_lightvae, mg_lightvae_v2
vae_device: "cuda:2"
checkpoints:
generator_ckpt: /path/to/longlive2/generator_checkpoint.pt
lora_ckpt: /path/to/longlive2/lora_checkpoint.pt
# Presence of this section enables LoRA inference; remove it for full-generator inference.
adapter:
type: lora
rank: 128
alpha: 128
dropout: 0.0
verbose: true
logging:
seed: 0
configs/inference_sp.yaml
+51
View File
@@ -0,0 +1,51 @@
# Ulysses sequence-parallel inference config for Wan2.2-TI2V-5B.
#
# Example launch:
# torchrun --nproc_per_node=4 inference_sp.py --config_path configs/inference_sp.yaml
model_kwargs:
model_name: Wan2.2-TI2V-5B
timestep_shift: 5.0
num_frame_per_block: 8
local_attn_size: 32
sp_size: 4
dp_size: 1
auto_sp_remainder: false
model_num_heads: 24
use_ema: false
output_folder: videos/longlive2_sp
num_samples: 1
save_latents_only: false
data:
data_path: /path/to/longlive2/inference_prompts
image_or_video_shape:
- 1
- 128
- 48
- 44
- 80
inference:
sampling_steps: 4
sink_size: 8
guidance_scale: 1.0
multi_shot_sink: true
multi_shot_rope_offset: 8
checkpoints:
generator_ckpt: /path/to/longlive2/generator_checkpoint.pt
lora_ckpt: /path/to/longlive2/lora_checkpoint.pt
# Presence of this section enables LoRA inference; remove it for full-generator inference.
adapter:
type: lora
rank: 128
alpha: 128
dropout: 0.0
verbose: true
logging:
seed: 0
configs/nvfp4/inference_i2v_nvfp4.yaml
+94
View File
@@ -0,0 +1,94 @@
# NVFP4 i2v inference with optional KV cache quantization and streaming VAE.
# = configs/nvfp4/inference_nvfp4.yaml + i2v conditioning.
# The first latent frame is encoded from the source image (provided by the
# video dataset), kept clean during denoising, and the rollout continues from it.
model_kwargs:
model_name: Wan2.2-TI2V-5B
timestep_shift: 5.0
num_frame_per_block: 8
local_attn_size: 32
sink_size: 8
# Enable i2v: build the video dataset, encode the first frame to a clean latent,
# and clamp it as the first-chunk conditioning during inference.
i2v: true
use_ema: false
output_folder: videos/longlive2_i2v_nvfp4
num_samples: 1
save_latents_only: false
save_with_index: true
inference_iter: -1
num_output_frames: 384
merge_lora: false
# i2v dataset / first-frame handling (read by MultiVideoConcatDataset in inference.py).
allow_padding: false
min_latent_frames: 0
max_chunks_per_shot: 0
uniform_prompt: false
data:
# For i2v the data_path points at a video/image dataset; the first frame of
# each clip is used as the conditioning image.
data_path: /path/to/longlive2/inference_video_dataset
image_or_video_shape:
- 1
- 384
- 48
- 44
- 80
inference:
sampling_steps: 4
# i2v conditioning: keep the first latent frame independent and clean.
independent_first_frame: true
# i2v training-time eval uses sink_size: 0; sink_size: 8 here matches the
# long-context streaming NVFP4 t2v setup. Tune as needed.
sink_size: 8
guidance_scale: 1.0
multi_shot_sink: true
multi_shot_rope_offset: 8
kv_quant: true
kv_quant_scale_rule: mse
kv_quant_backend: cuda
streaming_vae: false
async_vae: false
vae_type: wan # wan, mg_lightvae, mg_lightvae_v2
vae_device: "cuda:2"
negative_prompt: '色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走'
checkpoints:
generator_ckpt: /path/to/model_te.pt
# generator_ckpt: /path/to/model_4o6.pt # only set this if you are using FourOverSix (ckpt) for NVFP4
model_quant: true
model_quant_use_transformer_engine: true # only set this to true if you are using Transformer Engine (ckpt) for NVFP4; otherwise, set it to false (4o6)
model_quant_te_inference_only: true
model_quant_te_low_precision_weights: true
model_quant_te_fallback_to_fouroversix: false # only set this to false if you are using FourOverSix (ckpt) for NVFP4; otherwise, set it to true (TE)
model_quant_scale_rule: mse
model_quant_activation_scale_rule: mse
model_quant_weight_scale_rule: mse
model_quant_gradient_scale_rule: mse
torch_compile: auto
torch_compile_min_samples: 2
torch_compile_target: generator_model
torch_compile_backend: inductor
torch_compile_mode: max-autotune-no-cudagraphs
torch_compile_fullgraph: false
torch_compile_dynamic: false
torch_compile_suppress_errors: true
adapter:
type: lora
rank: 128
alpha: 128
dropout: 0.0
dtype: bfloat16
apply_to_critic: true
verbose: true
logging:
seed: 0
configs/nvfp4/inference_nvfp4.yaml
+75
View File
@@ -0,0 +1,75 @@
# NVFP4 inference with optional KV cache quantization and streaming VAE.
model_kwargs:
model_name: Wan2.2-TI2V-5B
timestep_shift: 5.0
num_frame_per_block: 8
local_attn_size: 32
sink_size: 8
use_ema: false
output_folder: videos/longlive2_nvfp4
num_samples: 1
save_latents_only: false
save_with_index: true
inference_iter: -1
num_output_frames: 384
merge_lora: false
data:
data_path: /path/to/longlive2/inference_prompts
image_or_video_shape:
- 1
- 384
- 48
- 44
- 80
inference:
sampling_steps: 4
sink_size: 8
guidance_scale: 1.0
multi_shot_sink: true
multi_shot_rope_offset: 8
kv_quant: true
kv_quant_scale_rule: mse
kv_quant_backend: cuda
streaming_vae: false
async_vae: false
vae_type: wan # wan, mg_lightvae, mg_lightvae_v2
vae_device: "cuda:2"
negative_prompt: '色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走'
checkpoints:
generator_ckpt: /path/to/model_te.pt
# generator_ckpt: /path/to/model_4o6.pt # only set this if you are using FourOverSix (ckpt) for NVFP4
model_quant: true
model_quant_use_transformer_engine: true # only set this to true if you are using Transformer Engine (ckpt) for NVFP4; otherwise, set it to false (4o6)
model_quant_te_inference_only: true
model_quant_te_low_precision_weights: true
model_quant_te_fallback_to_fouroversix: false # only set this to false if you are using FourOverSix (ckpt) for NVFP4; otherwise, set it to true (TE)
model_quant_scale_rule: mse
model_quant_activation_scale_rule: mse
model_quant_weight_scale_rule: mse
model_quant_gradient_scale_rule: mse
torch_compile: auto
torch_compile_min_samples: 2
torch_compile_target: generator_model
torch_compile_backend: inductor
torch_compile_mode: max-autotune-no-cudagraphs
torch_compile_fullgraph: false
torch_compile_dynamic: false
torch_compile_suppress_errors: true
adapter:
type: lora
rank: 128
alpha: 128
dropout: 0.0
dtype: bfloat16
apply_to_critic: true
verbose: true
logging:
seed: 0
configs/nvfp4/train_ar_nvfp4.yaml
+96
View File
@@ -0,0 +1,96 @@
# NVFP4 AR teacher-forcing training.
infra:
sequence_parallel_size: 4
sharding_strategy: hybrid_full
mixed_precision: true
gradient_checkpointing: true
vae_halo_latents: 28
generator_fsdp_wrap_strategy: size
text_encoder_fsdp_wrap_strategy: size
model_quant: true
model_quant_scale_rule: static_6
model_quant_activation_scale_rule: mse
model_quant_weight_scale_rule: static_6
model_quant_gradient_scale_rule: mse
model_kwargs:
model_name: Wan2.2-TI2V-5B
timestep_shift: 5.0
num_frame_per_block: 8
local_attn_size: -1
sink_size: 8
t_scale: 1.0
rope_method: linear
checkpoints:
generator_ckpt: /path/to/wan_or_ar_checkpoint.pt
algorithm:
trainer: diffusion
causal: true
teacher_forcing: true
num_train_timestep: 1000
denoising_loss_type: flow
training:
lr: 1.0e-05
weight_decay: 0.0
beta1: 0.0
beta2: 0.999
batch_size: 1
gradient_accumulation_steps: 4
ema_weight: 0.99
ema_start_step: 100
log_iters: 15
max_checkpoints: 10
max_iters: 600
gc_interval: 100
error_recycling:
enabled: true
num_buckets: 50
buffer_size_per_bucket: 500
buffer_warmup_iter: 50
context_inject_prob: 0.9
latent_inject_prob: 0.9
noise_inject_prob: 0.01
clean_prob: 0.2
clean_buffer_update_prob: 0.1
modulate_factor: 0.3
start_step: 0
replacement_strategy: random
data:
data_path: /path/to/longlive2/train_dataset_64s
eval_data_path: /path/to/longlive2/eval_prompts
image_or_video_shape:
- 1
- 384
- 48
- 44
- 80
load_raw_video: true
max_chunks_per_shot: 4
inference:
sampling_steps: 50
sink_size: 8
local_attn_size: 24
multi_shot_sink: true
multi_shot_rope_offset: 8
guidance_scale: 3.0
use_relative_rope: false
inference_t_scale: 1.0
negative_prompt: '色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走'
evaluation:
interval: 1000
num_frames: 384
use_ema: true
val_batch_size: 1
save_latents_only: true
logging:
seed: 0
wandb_key: null
wandb_entity: null
wandb_project: LongLive2-AR-NVFP4
configs/nvfp4/train_dmd_nvfp4_step4.yaml
+126
View File
@@ -0,0 +1,126 @@
# NVFP4 DMD LoRA distillation, 4-step rollout.
infra:
sharding_strategy: hybrid_full
mixed_precision: true
gradient_checkpointing: true
generator_fsdp_wrap_strategy: size
real_score_fsdp_wrap_strategy: size
fake_score_fsdp_wrap_strategy: size
text_encoder_fsdp_wrap_strategy: size
generator_quant: true
real_score_quant: true
fake_score_quant: false
real_score_quant_materialize: true
generator_quant_scale_rule: mse
generator_quant_activation_scale_rule: mse
generator_quant_weight_scale_rule: mse
generator_quant_gradient_scale_rule: mse
generator_quant_use_default_filtered_modules: false
real_score_quant_scale_rule: mse
real_score_quant_activation_scale_rule: mse
real_score_quant_weight_scale_rule: mse
real_score_quant_use_default_filtered_modules: false
generator_quant_filtered_modules: &nvfp4_filtered_modules
- text_embedding.0
- text_embedding.2
- patch_embedding
- time_projection.1
- time_embedding.0
- time_embedding.2
- head.head
- head.modulation
- re:.*norm_k$
- re:.*norm_q$
- re:.*norm1$
- re:.*norm2$
- re:.*norm3$
real_score_quant_filtered_modules: *nvfp4_filtered_modules
model_kwargs:
model_name: Wan2.2-TI2V-5B
timestep_shift: 5.0
num_frame_per_block: 8
local_attn_size: 32
sink_size: 8
checkpoints:
generator_ckpt: null
real_score_ckpt: null
fake_score_ckpt: null
algorithm:
trainer: score_distillation
distribution_loss: dmd
all_causal: true
ts_schedule: false
num_train_timestep: 1000
denoising_loss_type: flow
warp_denoising_step: false
denoising_step_list: [1000, 946, 854, 681, 0]
teacher_forcing: false
backward_simulation: true
real_guidance_scale: 3.0
fake_guidance_scale: 0.0
training:
lr: 1.0e-05
lr_critic: 2.0e-06
weight_decay: 0.0
beta1: 0.0
beta2: 0.999
beta1_critic: 0.0
beta2_critic: 0.999
batch_size: 2
gradient_accumulation_steps: 1
ema_weight: 0.99
ema_start_step: 200
log_iters: 50
max_checkpoints: 20
max_iters: 5000
gc_interval: 100
dfake_gen_update_ratio: 5
min_num_training_frames: 32
num_training_frames: 32
slice_last_frames: 32
chunks_per_shot: 2
data:
data_path: /path/to/longlive2/train_dataset_16to32s
eval_data_path: /path/to/longlive2/eval_prompts
image_or_video_shape:
- 1
- 32
- 48
- 44
- 80
load_raw_video: true
inference:
sampling_steps: 4
sink_size: 8
multi_shot_sink: true
multi_shot_rope_offset: 8
narrative_rope_phase_shift: 8
negative_prompt: '色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走'
evaluation:
interval: 99999999999
num_frames: 32
use_ema: false
val_batch_size: 1
save_latents_only: true
adapter:
type: lora
rank: 128
alpha: 128
dropout: 0.0
dtype: bfloat16
apply_to_critic: true
verbose: true
logging:
seed: 0
wandb_key: null
wandb_entity: null
wandb_project: LongLive2-DMD-NVFP4
configs/nvfp4/train_i2v_ar_nvfp4.yaml
+101
View File
@@ -0,0 +1,101 @@
# NVFP4 i2v AR teacher-forcing training.
# = configs/train_i2v_ar.yaml + the NVFP4 quant infra from
# configs/nvfp4/train_ar_nvfp4.yaml.
infra:
sequence_parallel_size: 4
sharding_strategy: hybrid_full
mixed_precision: true
gradient_checkpointing: true
vae_halo_latents: 28
generator_fsdp_wrap_strategy: size
text_encoder_fsdp_wrap_strategy: size
model_quant: true
model_quant_scale_rule: static_6
model_quant_activation_scale_rule: mse
model_quant_weight_scale_rule: static_6
model_quant_gradient_scale_rule: mse
model_kwargs:
model_name: Wan2.2-TI2V-5B
timestep_shift: 5.0
num_frame_per_block: 8
local_attn_size: -1
sink_size: 8
t_scale: 1.0
rope_method: linear
checkpoints:
generator_ckpt: /path/to/wan_or_ar_checkpoint.pt
# I2V AR diffusion objective. The first latent frame is encoded from the
# source image, kept clean, and masked out of the flow-matching loss.
algorithm:
trainer: diffusion
i2v: true
causal: true
teacher_forcing: true
independent_first_frame: true
num_train_timestep: 1000
denoising_loss_type: flow
training:
lr: 1.0e-05
weight_decay: 0.0
beta1: 0.0
beta2: 0.999
batch_size: 1
gradient_accumulation_steps: 1
ema_weight: 0.99
ema_start_step: 100
log_iters: 15
max_checkpoints: 20
max_iters: 600
gc_interval: 100
error_recycling:
enabled: true
num_buckets: 50
buffer_size_per_bucket: 500
buffer_warmup_iter: 50
context_inject_prob: 0.9
latent_inject_prob: 0.9
noise_inject_prob: 0.01
clean_prob: 0.2
clean_buffer_update_prob: 0.1
modulate_factor: 0.3
start_step: 0
replacement_strategy: random
data:
data_path: /path/to/longlive2/train_video_dataset
eval_data_path: /path/to/longlive2/eval_video_dataset
image_or_video_shape:
- 1
- 96
- 48
- 44
- 80
load_raw_video: true
inference:
sampling_steps: 50
sink_size: 0
multi_shot_sink: true
multi_shot_rope_offset: 8
guidance_scale: 3.0
independent_first_frame: true
use_relative_rope: false
inference_t_scale: 1.0
negative_prompt: '色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走'
evaluation:
interval: 1000
num_frames: 32
use_ema: true
val_batch_size: 1
save_latents_only: true
logging:
seed: 0
wandb_key: null
wandb_entity: null
wandb_project: LongLive2-I2V-AR-NVFP4
configs/nvfp4/train_i2v_dmd_nvfp4_step4.yaml
+142
View File
@@ -0,0 +1,142 @@
# NVFP4 i2v DMD LoRA distillation, 4-step rollout.
# = configs/train_i2v_dmd.yaml + the NVFP4 quant infra and 4-step denoising
# schedule from configs/nvfp4/train_dmd_nvfp4_step4.yaml.
infra:
sharding_strategy: hybrid_full
mixed_precision: true
gradient_checkpointing: true
generator_fsdp_wrap_strategy: size
real_score_fsdp_wrap_strategy: size
fake_score_fsdp_wrap_strategy: size
text_encoder_fsdp_wrap_strategy: size
generator_quant: true
real_score_quant: true
fake_score_quant: false
real_score_quant_materialize: true
generator_quant_scale_rule: mse
generator_quant_activation_scale_rule: mse
generator_quant_weight_scale_rule: mse
generator_quant_gradient_scale_rule: mse
generator_quant_use_default_filtered_modules: false
real_score_quant_scale_rule: mse
real_score_quant_activation_scale_rule: mse
real_score_quant_weight_scale_rule: mse
real_score_quant_use_default_filtered_modules: false
generator_quant_filtered_modules: &nvfp4_filtered_modules
- text_embedding.0
- text_embedding.2
- patch_embedding
- time_projection.1
- time_embedding.0
- time_embedding.2
- head.head
- head.modulation
- re:.*norm_k$
- re:.*norm_q$
- re:.*norm1$
- re:.*norm2$
- re:.*norm3$
real_score_quant_filtered_modules: *nvfp4_filtered_modules
model_kwargs:
model_name: Wan2.2-TI2V-5B
timestep_shift: 5.0
num_frame_per_block: 8
local_attn_size: 32
sink_size: 8
# Optional initialization checkpoints. In practice, set generator_ckpt to the
# i2v AR checkpoint used as the DMD student initialization.
checkpoints:
generator_ckpt: null
real_score_ckpt: null
fake_score_ckpt: null
# I2V DMD objective. Backward simulation uses the same first-chunk image
# conditioning rule as i2v inference: the clean first latent is injected during
# denoising and excluded from generator/critic losses.
algorithm:
trainer: score_distillation
distribution_loss: dmd
i2v: true
all_causal: true
ts_schedule: false
num_train_timestep: 1000
denoising_loss_type: flow
warp_denoising_step: false
denoising_step_list: [1000, 946, 854, 681, 0]
teacher_forcing: false
backward_simulation: true
independent_first_frame: true
real_guidance_scale: 3.0
fake_guidance_scale: 0.0
training:
lr: 1.0e-05
lr_critic: 2.0e-06
weight_decay: 0.0
beta1: 0.0
beta2: 0.999
beta1_critic: 0.0
beta2_critic: 0.999
batch_size: 1
gradient_accumulation_steps: 1
ema_weight: 0.99
ema_start_step: 200
log_iters: 50
max_checkpoints: 50
max_iters: 5000
gc_interval: 100
dfake_gen_update_ratio: 5
# Backward simulation samples a rollout length in [min_num_training_frames, num_training_frames].
# Setting both to 32 gives a fixed 32-latent DMD rollout.
min_num_training_frames: 32
num_training_frames: 32
# Keep this many tail latents for the DMD/critic losses after backward simulation.
# With a 32-latent rollout, 32 keeps the full generated window.
slice_last_frames: 32
chunks_per_shot: 2
data:
data_path: /path/to/longlive2/train_video_dataset
eval_data_path: /path/to/longlive2/eval_video_dataset
image_or_video_shape:
- 1
- 32
- 48
- 44
- 80
load_raw_video: true
# The distillation loop generates i2v trajectories, then computes DMD losses.
inference:
sampling_steps: 4
sink_size: 0
multi_shot_sink: true
multi_shot_rope_offset: 8
independent_first_frame: true
negative_prompt: '色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走'
# Optional validation generation from the distilled model. Uses inference sampling/cache settings by default.
evaluation:
interval: 99999999999
num_frames: 32
use_ema: false
val_batch_size: 1
save_latents_only: true
# Presence of this section enables LoRA distillation; remove it for full fine-tuning.
adapter:
type: lora
rank: 128
alpha: 128
dropout: 0.0
dtype: bfloat16
apply_to_critic: true
verbose: true
logging:
seed: 0
wandb_key: null
wandb_entity: null
wandb_project: LongLive2-I2V-DMD-NVFP4
configs/train_ar.yaml
+85
View File
@@ -0,0 +1,85 @@
# Distributed training runtime.
infra:
sequence_parallel_size: 4
sharding_strategy: hybrid_full
mixed_precision: true
gradient_checkpointing: true
vae_halo_latents: 28
generator_fsdp_wrap_strategy: size
text_encoder_fsdp_wrap_strategy: size
# Passed directly to WanDiffusionWrapper(**model_kwargs).
model_kwargs:
model_name: Wan2.2-TI2V-5B
timestep_shift: 5.0
num_frame_per_block: 8
local_attn_size: -1
# AR diffusion objective.
algorithm:
trainer: diffusion
causal: true
teacher_forcing: true
num_train_timestep: 1000
denoising_loss_type: flow
training:
lr: 1.0e-05
weight_decay: 0.0
beta1: 0.0
beta2: 0.999
batch_size: 1
gradient_accumulation_steps: 1
ema_weight: 0.99
ema_start_step: 100
log_iters: 1
max_checkpoints: 20
max_iters: 600
gc_interval: 100
error_recycling:
enabled: true
num_buckets: 50
buffer_size_per_bucket: 500
buffer_warmup_iter: 50
context_inject_prob: 0.9
latent_inject_prob: 0.9
noise_inject_prob: 0.01
clean_prob: 0.2
clean_buffer_update_prob: 0.1
modulate_factor: 0.3
start_step: 0
replacement_strategy: random
data:
data_path: /path/to/longlive2/train_dataset
eval_data_path: /path/to/longlive2/eval_prompts
image_or_video_shape:
- 1
- 96
- 48
- 44
- 80
load_raw_video: true
# Generation settings used by training-time evaluation.
inference:
sampling_steps: 50
sink_size: 0
multi_shot_sink: true
multi_shot_rope_offset: 8
guidance_scale: 3.0
negative_prompt: '色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走'
# Lightweight validation generation during AR training.
evaluation:
interval: 1
num_frames: 16
use_ema: true
val_batch_size: 1
save_latents_only: true
logging:
seed: 0
wandb_key: null
wandb_entity: null
wandb_project: LongLive2-AR
configs/train_dmd.yaml
+102
View File
@@ -0,0 +1,102 @@
# Distributed distillation runtime.
infra:
sharding_strategy: hybrid_full
mixed_precision: true
gradient_checkpointing: true
generator_fsdp_wrap_strategy: size
real_score_fsdp_wrap_strategy: size
fake_score_fsdp_wrap_strategy: size
text_encoder_fsdp_wrap_strategy: size
# Shared Wan backbone settings for the student, real-score teacher, and fake-score critic.
# Role-specific model kwargs can still be provided as real_model_kwargs/fake_model_kwargs
# if a future run needs different model construction.
model_kwargs:
model_name: Wan2.2-TI2V-5B
timestep_shift: 5.0
num_frame_per_block: 8
local_attn_size: 32
# Optional initialization checkpoints. These are flattened by normalize_config
# so the trainer still reads generator_ckpt/real_score_ckpt/fake_score_ckpt.
# With adapter enabled, generator_ckpt initializes the base generator before LoRA wrapping.
checkpoints:
generator_ckpt: null
real_score_ckpt: null
fake_score_ckpt: null
# DMD distillation objective.
algorithm:
trainer: score_distillation
distribution_loss: dmd
all_causal: true
ts_schedule: false
real_guidance_scale: 3.0
fake_guidance_scale: 0.0
training:
lr: 1.0e-05
lr_critic: 2.0e-06
weight_decay: 0.0
beta1: 0.0
beta2: 0.999
beta1_critic: 0.0
beta2_critic: 0.999
batch_size: 1
gradient_accumulation_steps: 1
ema_weight: 0.99
ema_start_step: 200
log_iters: 50
max_checkpoints: 50
max_iters: 5000
gc_interval: 100
dfake_gen_update_ratio: 5
# Backward simulation samples a rollout length in [min_num_training_frames, num_training_frames].
# Setting both to 32 gives a fixed 32-latent DMD rollout.
min_num_training_frames: 32
num_training_frames: 32
# Keep this many tail latents for the DMD/critic losses after backward simulation.
# With a 32-latent rollout, 32 keeps the full generated window.
slice_last_frames: 32
data:
data_path: /path/to/longlive2/train_dataset
eval_data_path: /path/to/longlive2/eval_prompts
image_or_video_shape:
- 1
- 32
- 48
- 44
- 80
# The distillation loop generates video trajectories, then computes DMD losses.
inference:
sampling_steps: 4
sink_size: 0
multi_shot_rope_offset: 8
# Optional validation generation from the distilled model. Uses inference sampling/cache settings by default.
evaluation:
interval: 1
num_frames: 32
use_ema: false
# Can be raised when validation memory allows; 1 is the safe default.
val_batch_size: 1
save_latents_only: true
# Presence of this section enables LoRA distillation; remove it for full fine-tuning.
adapter:
type: lora
rank: 128
alpha: 128
dropout: 0.0
apply_to_critic: true
verbose: true
logging:
seed: 0
wandb_key: null
wandb_entity: null
wandb_project: LongLive2-DMD
configs/train_i2v_ar.yaml
+89
View File
@@ -0,0 +1,89 @@
# Distributed i2v AR training runtime.
infra:
sequence_parallel_size: 4
sharding_strategy: hybrid_full
mixed_precision: true
gradient_checkpointing: true
vae_halo_latents: 28
generator_fsdp_wrap_strategy: size
text_encoder_fsdp_wrap_strategy: size
# Passed directly to WanDiffusionWrapper(**model_kwargs).
model_kwargs:
model_name: Wan2.2-TI2V-5B
timestep_shift: 5.0
num_frame_per_block: 8
local_attn_size: -1
# I2V AR diffusion objective. The first latent frame is encoded from the
# source image, kept clean, and masked out of the flow-matching loss.
algorithm:
trainer: diffusion
i2v: true
causal: true
teacher_forcing: true
independent_first_frame: true
num_train_timestep: 1000
denoising_loss_type: flow
training:
lr: 1.0e-05
weight_decay: 0.0
beta1: 0.0
beta2: 0.999
batch_size: 1
gradient_accumulation_steps: 1
ema_weight: 0.99
ema_start_step: 100
log_iters: 1
max_checkpoints: 20
max_iters: 600
gc_interval: 100
error_recycling:
enabled: true
num_buckets: 50
buffer_size_per_bucket: 500
buffer_warmup_iter: 50
context_inject_prob: 0.9
latent_inject_prob: 0.9
noise_inject_prob: 0.01
clean_prob: 0.2
clean_buffer_update_prob: 0.1
modulate_factor: 0.3
start_step: 0
replacement_strategy: random
data:
data_path: /path/to/longlive2/train_video_dataset
eval_data_path: /path/to/longlive2/eval_video_dataset
image_or_video_shape:
- 1
- 96
- 48
- 44
- 80
load_raw_video: true
# Generation settings used by training-time evaluation.
inference:
sampling_steps: 50
sink_size: 0
multi_shot_sink: true
multi_shot_rope_offset: 8
guidance_scale: 3.0
independent_first_frame: true
negative_prompt: '色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走'
# Lightweight validation generation during i2v AR training.
evaluation:
interval: 1
num_frames: 32
use_ema: true
val_batch_size: 1
save_latents_only: true
logging:
seed: 0
wandb_key: null
wandb_entity: null
wandb_project: LongLive2-I2V-AR
configs/train_i2v_dmd.yaml
+104
View File
@@ -0,0 +1,104 @@
# Distributed i2v DMD distillation runtime.
infra:
sharding_strategy: hybrid_full
mixed_precision: true
gradient_checkpointing: true
generator_fsdp_wrap_strategy: size
real_score_fsdp_wrap_strategy: size
fake_score_fsdp_wrap_strategy: size
text_encoder_fsdp_wrap_strategy: size
# Shared Wan backbone settings for the student, real-score teacher, and fake-score critic.
model_kwargs:
model_name: Wan2.2-TI2V-5B
timestep_shift: 5.0
num_frame_per_block: 8
local_attn_size: 32
# Optional initialization checkpoints. In practice, set generator_ckpt to the
# i2v AR checkpoint used as the DMD student initialization.
checkpoints:
generator_ckpt: null
real_score_ckpt: null
fake_score_ckpt: null
# I2V DMD objective. Backward simulation uses the same first-chunk image
# conditioning rule as i2v inference: the clean first latent is injected during
# denoising and excluded from generator/critic losses.
algorithm:
trainer: score_distillation
distribution_loss: dmd
i2v: true
all_causal: true
backward_simulation: true
teacher_forcing: false
independent_first_frame: true
ts_schedule: false
real_guidance_scale: 3.0
fake_guidance_scale: 0.0
training:
lr: 1.0e-05
lr_critic: 2.0e-06
weight_decay: 0.0
beta1: 0.0
beta2: 0.999
beta1_critic: 0.0
beta2_critic: 0.999
batch_size: 1
gradient_accumulation_steps: 1
ema_weight: 0.99
ema_start_step: 200
log_iters: 50
max_checkpoints: 50
max_iters: 5000
gc_interval: 100
dfake_gen_update_ratio: 5
# Backward simulation samples a rollout length in [min_num_training_frames, num_training_frames].
# Setting both to 32 gives a fixed 32-latent DMD rollout.
min_num_training_frames: 32
num_training_frames: 32
# Keep this many tail latents for the DMD/critic losses after backward simulation.
# With a 32-latent rollout, 32 keeps the full generated window.
slice_last_frames: 32
data:
data_path: /path/to/longlive2/train_video_dataset
eval_data_path: /path/to/longlive2/eval_video_dataset
image_or_video_shape:
- 1
- 32
- 48
- 44
- 80
load_raw_video: true
# The distillation loop generates i2v trajectories, then computes DMD losses.
inference:
sampling_steps: 4
sink_size: 0
multi_shot_rope_offset: 8
independent_first_frame: true
# Optional validation generation from the distilled model. Uses inference sampling/cache settings by default.
evaluation:
interval: 1
num_frames: 32
use_ema: false
val_batch_size: 1
save_latents_only: true
# Presence of this section enables LoRA distillation; remove it for full fine-tuning.
adapter:
type: lora
rank: 128
alpha: 128
dropout: 0.0
apply_to_critic: true
verbose: true
logging:
seed: 0
wandb_key: null
wandb_entity: null
wandb_project: LongLive2-I2V-DMD
docs/FLASH_ATTENTION_3_AND_HOPPER_SUPPORT.md
+103
View File
@@ -0,0 +1,103 @@
# Flash Attention 3 and Hopper GPU Support
This document describes the Flash Attention 3 (FA3) integration and extended Hopper GPU support in LongLive.
## Overview
LongLive supports both Flash Attention 2 (FA2) and Flash Attention 3 (FA3) for efficient attention computation. FA3 is automatically enabled on Hopper architecture GPUs (Compute Capability 9.0+), providing improved performance.
## Supported Hardware
### Hopper Architecture GPUs (FA3 Enabled)
- **NVIDIA H100** - Data center GPU
- **NVIDIA H800** - China-specific variant
- **NVIDIA H20** - China-specific variant
All Hopper GPUs share Compute Capability 9.0, which is the requirement for FA3.
### Other GPUs (FA2 Fallback)
- **NVIDIA A100** - Ampere architecture (Compute Capability 8.0)
- **NVIDIA A800** - Ampere architecture (Compute Capability 8.0)
- Other CUDA-capable GPUs with FA2 support
## Design Choices
### 1. GPU Detection via Compute Capability
Instead of relying on device name string matching (which would miss H800/H20), we detect Hopper GPUs using CUDA Compute Capability:
```python
def is_hopper_gpu():
if torch.cuda.is_available():
major, _ = torch.cuda.get_device_capability()
return major >= 9 # Hopper Compute Capability == 9.0
return False
```
**Rationale:**
- Device names vary across vendors and regions (H100, H800, H20, etc.)
- Compute Capability is a reliable, standardized way to identify GPU architecture
- All Hopper GPUs report `major=9` regardless of their marketing name
### 2. FA3 Return Value Handling
Flash Attention 3's `flash_attn_varlen_func` has a different return signature than FA2:
| Version | Return Value |
|---------|--------------|
| FA2 | `(output, softmax_lse, ...)` - tuple, use `[0]` to get output |
| FA3 | `output` - tensor directly |
The code correctly handles this difference:
```python
# FA3 path - direct tensor return
x = flash_attn_interface.flash_attn_varlen_func(...).unflatten(0, (b, lq))
# FA2 path - tuple return (handled in else branch)
x = flash_attn.flash_attn_varlen_func(...).unflatten(0, (b, lq))
```
### 3. Automatic Fallback
The system gracefully falls back to FA2 when FA3 is unavailable:
- If `flash_attn_interface` module is not installed
- If running on non-Hopper GPU
- If user explicitly requests FA2 via `version=2` parameter
A warning is issued when FA3 is explicitly requested but unavailable.
## Usage
### Automatic Selection (Recommended)
By default, LongLive automatically selects the optimal attention implementation:
```python
from wan_5b.modules.attention import attention
# FA3 will be used on Hopper GPUs, FA2 otherwise
output = attention(q, k, v)
```
### Explicit Version Selection
You can force a specific Flash Attention version:
```python
# Force FA2 (useful for debugging or compatibility)
output = attention(q, k, v, fa_version=2)
# Request FA3 (falls back to FA2 with warning if unavailable)
output = attention(q, k, v, fa_version=3)
```
## Installation
### Flash Attention 3 (Hopper GPUs)
```bash
git clone https://github.com/Dao-AILab/flash-attention.git
cd flash-attention/hopper
python setup.py install
```
docs/getting_started.md
+430
View File
@@ -0,0 +1,430 @@
# LongLive2.0 Usage
This document contains the release commands for installation, training, inference, and utilities. The root README keeps the project overview and paper figures.
## Installation
Create a Python 3.10 environment and install the required packages:
```bash
conda create -n longlive2 python=3.10 -y
conda activate longlive2
pip install torch==2.8.0 torchvision==0.23.0 --index-url https://download.pytorch.org/whl/cu128
pip install -r requirements.txt
pip install flash-attn --no-build-isolation
```
TensorRT is not required for the default training or inference path. If a
TensorRT utility is needed, install it separately after the base requirements:
```bash
pip install nvidia-pyindex
pip install nvidia-tensorrt
pip install pycuda
```
Download the Wan2.2-TI2V-5B components and replace the `/path/to/longlive2.0/...` placeholders in the config files before running training or inference.
If you set `inference.vae_type` to `mg_lightvae` or `mg_lightvae_v2`, download
the corresponding VAE checkpoints from the Hugging Face repository
`Skywork/Matrix-Game-3.0` and place them under `wan_models/Matrix-Game-3.0/`:
```text
wan_models/Matrix-Game-3.0/MG-LightVAE.pth
wan_models/Matrix-Game-3.0/MG-LightVAE_v2.pth
```
### NVFP4 Environment
The default installation above is the clean BF16 release setup. NVFP4 training
and inference use local CUDA extensions and are more version-sensitive, so keep
them in a separate environment.
Known-good NVFP4 baseline inherited from the Sage branch:
```text
Python: 3.12.12
PyTorch: 2.10.0+cu128
TorchVision: 0.25.0+cu128
CUDA target: 12.8
FlashAttention: 2.8.3, built from source
```
Create or activate the NVFP4 environment:
```bash
conda create -n longlive2_nvfp4 python=3.12 -y
conda activate longlive2_nvfp4
conda install -c nvidia cuda-toolkit=12.8 -y
pip install -r requirements.txt
pip install --upgrade --index-url https://download.pytorch.org/whl/cu128 \
torch==2.10.0 torchvision==0.25.0
pip install --upgrade torchao==0.16.0
```
If you already have a working `qlive` environment from LongLive_Sage, you can
activate it instead of creating `longlive2_nvfp4`.
Verify the Torch/CUDA pair:
```bash
python -c "import torch, torchvision; print(torch.__version__, torch.version.cuda); print(torchvision.__version__)"
```
Build the modified local `fouroversix` package:
```bash
cd fouroversix
pip install ninja packaging psutil "setuptools>=77.0.3"
# Optional: limit compile targets.
export CUDA_ARCHS=100 # B200 / GB200 / GB300
# export CUDA_ARCHS=120 # RTX 50/60 series, if needed
pip install --no-build-isolation -e .
cd ..
```
Build FlashAttention from source, rather than relying on a prebuilt wheel:
```bash
git clone https://github.com/Dao-AILab/flash-attention.git
cd flash-attention
git checkout v2.8.3
pip install -U pip setuptools wheel ninja packaging
pip install --no-build-isolation -e .
cd ..
```
Install TransformerEngine if `model_quant_use_transformer_engine: true` will be
used:
```bash
python -m pip install --no-build-isolation "transformer-engine[pytorch]"
```
Build the fused LongLive FP4 KV-cache dequant extension:
```bash
cd utils/kernel
python setup.py build_ext --inplace
cd ../..
```
Quick NVFP4 checks:
```bash
python -c "import flash_attn; print(flash_attn.__version__)"
python -c "import fouroversix; from utils.quant import LongLiveQuantizationConfig, quantize_to_fp4"
python -c "from utils.kernel.kv_dequant import dequantize_kv_cache_fp4"
```
The release NVFP4 configs and direct run commands are summarized below. See
`README_NVFP4.md` for lower-level implementation notes.
## Configs
The release keeps three main configs:
```text
configs/train_ar.yaml # AR diffusion training
configs/train_dmd.yaml # DMD distillation
configs/inference.yaml # inference
```
TorchAO FP8 PTQ inference has a separate config:
```text
configs/fp8/inference_fp8.yaml
```
The NVFP4 path keeps its configs separate from the default BF16 release path:
```text
configs/nvfp4/train_ar_nvfp4.yaml # stage 1 AR teacher-forcing training
configs/nvfp4/train_dmd_nvfp4_step4.yaml # stage 2 DMD LoRA distillation, 4-step rollout
configs/nvfp4/inference_nvfp4.yaml # NVFP4 inference with optional KV quantization
```
The configs use a shared organization:
- `model_kwargs`: arguments passed to `WanDiffusionWrapper`.
- `infra`: distributed training/runtime settings.
- `algorithm`: AR or DMD objective settings.
- `training`: optimizer, batch size, checkpoint cadence, and loop settings.
- `data`: training or prompt data paths.
- `inference`: sampling and cache settings.
- `checkpoints`: model and LoRA checkpoint paths.
- `adapter`: optional LoRA settings. Remove this section to disable LoRA.
## Training
### AR Diffusion Training
Edit `configs/train_ar.yaml` to set the dataset path, evaluation prompt path, logging path, and distributed runtime settings. Then run:
```bash
torchrun --standalone --nnodes=1 --nproc_per_node=8 train.py \
--config_path configs/train_ar.yaml \
--logdir logs/test_train_ar \
--wandb-save-dir wandb \
--disable-wandb
```
Notes:
- `infra.sequence_parallel_size` controls the SP group size.
- `infra.vae_halo_latents` controls chunk-halo VAE preparation.
- `model_kwargs.local_attn_size` is a model construction setting.
- `inference.sink_size`, `inference.multi_shot_sink`, and `inference.multi_shot_rope_offset` control evaluation-time generation during training.
### DMD Distillation
Edit `configs/train_dmd.yaml` to set the dataset path and initialization checkpoints. Then run:
```bash
torchrun --standalone --nnodes=1 --nproc_per_node=8 train.py \
--config_path configs/train_dmd.yaml \
--logdir logs/test_train_dmd \
--wandb-save-dir wandb \
--disable-wandb
```
Notes:
- `algorithm.real_guidance_scale` and `algorithm.fake_guidance_scale` are used by score distillation.
- `inference.sampling_steps` controls the distillation rollout sampling steps.
- If `adapter` is present, LoRA distillation is enabled. Otherwise the generator is fully fine-tuned.
- Auto-resume is enabled by default unless `--no-auto-resume` is passed.
### NVFP4 Training
Use the `longlive2_nvfp4` environment and build the NVFP4 extensions before
running these commands. Replace the `/path/to/...` placeholders in the configs
first.
Stage 1 trains the NVFP4 AR teacher-forcing model:
```bash
torchrun --standalone --nnodes=1 --nproc_per_node=4 train.py \
--config_path configs/nvfp4/train_ar_nvfp4.yaml \
--logdir logs/nvfp4_ar \
--wandb-save-dir wandb \
--disable-wandb
```
Stage 2 runs NVFP4 DMD LoRA distillation from the AR checkpoint:
```bash
torchrun --standalone --nnodes=1 --nproc_per_node=4 train.py \
--config_path configs/nvfp4/train_dmd_nvfp4_step4.yaml \
--logdir logs/nvfp4_dmd_step4 \
--wandb-save-dir wandb \
--disable-wandb
```
Notes:
- `--nproc_per_node` controls the per-node GPU count. The NVFP4 examples use 4
GPUs; set it to 8 or another value for your machine.
- `infra.model_quant` enables NVFP4 generator training for stage 1.
- `infra.generator_quant`, `infra.real_score_quant`, and
`infra.fake_score_quant` choose which DMD networks use NVFP4 in stage 2.
After stage 1 and stage 2 are complete, you can pre-merge the AR generator and
DMD LoRA weights for inference. The export script reads `generator_ckpt`,
`lora_ckpt`, `adapter`, and `model_quant_*` from the NVFP4 inference config.
To save a compact FourOverSix materialized NVFP4 generator checkpoint:
```bash
python scripts/save_merged_nvfp4_generator.py \
--config_path configs/nvfp4/inference_nvfp4.yaml \
--output_path /path/to/model_4o6.pt \
--backend fouroversix \
--device cuda:0
```
To save merged BF16 weights for TransformerEngine runtime quantization:
```bash
python scripts/save_merged_nvfp4_generator.py \
--config_path configs/nvfp4/inference_nvfp4.yaml \
--output_path /path/to/model_te.pt \
--backend transformer_engine \
--device cuda:0
```
The `fouroversix` export is the small packed/materialized NVFP4 checkpoint. The
`transformer_engine` export intentionally saves merged BF16 weights, because a
TransformerEngine module `state_dict` is not a compact packed NVFP4 storage
format; TE quantization is applied again when inference loads the BF16 weights.
### Merge Generator and LoRA Weights
For the regular BF16 release path, you can pre-merge the AR generator checkpoint
and DMD LoRA checkpoint into one reusable generator checkpoint. This keeps
quick-start inference simple: inference only loads `checkpoints.generator_ckpt`
and does not need to construct or load LoRA adapters at runtime.
```bash
python scripts/merge_lora_generator.py \
--config_path configs/inference.yaml \
--output_path /path/to/longlive2_merged_generator.pt \
--device cuda:0
```
After the merge, set `checkpoints.generator_ckpt` in `configs/inference.yaml` to
the merged checkpoint. If you run the full `inference.py` entry point, remove or
unset `checkpoints.lora_ckpt` and the `adapter` section so LoRA is not applied a
second time.
## Inference
Edit `configs/inference.yaml` to set:
- `data.data_path`: prompt folder.
- `checkpoints.generator_ckpt`: merged generator checkpoint.
- `output_folder`: output video directory.
- `num_samples`: number of sampled videos per prompt.
Run:
```bash
torchrun --standalone --nnodes=1 --nproc_per_node=8 inference.py \
--config_path configs/inference.yaml
```
Inference notes:
- `inference.sampling_steps` controls the number of denoising steps.
- `inference.guidance_scale` controls inference CFG.
- `inference.sink_size` controls the standard attention sink size.
- `inference.multi_shot_sink` enables the multi-shot attention sink.
- `inference.multi_shot_rope_offset` controls the multi-shot RoPE offset.
### FP8 PTQ Inference
Set `checkpoints.generator_ckpt` in `configs/fp8/inference_fp8.yaml` to the
downloaded merged BF16 `model_bf16.pt`, then run:
```bash
python inference.py --config_path configs/fp8/inference_fp8.yaml
```
`fp8_quant: true` applies TorchAO row-wise dynamic W8A8 quantization after the
generator has been loaded and converted to BF16, and before `torch.compile`.
It cannot be combined with `model_quant: true`, which selects the NVFP4 path.
With the provided 5B model, 300 eligible core Linear layers use FP8 while six
small conditioning/output projections remain BF16 for stability and to avoid
FP8 overhead.
The validated stack is Python 3.10, PyTorch 2.8.0+cu128, and TorchAO 0.13.0 on
H100 (SM90); compute capability 8.9 or newer is required. The supplied config
uses `torch_compile: auto`: it skips compilation when `inference_iter`
explicitly limits the run to fewer than three samples, and enables it when all
prompts are requested. Its `max-autotune` warm-up can take several minutes
while guard/shape variants are compiled. Use repeated inference and discard all
compile/warm-up samples when measuring steady-state performance; set
`torch_compile: false` for a short eager-mode smoke test.
The supplied config uses the single 8-latent-frame block validated on H100.
Longer generation introduces additional KV-cache shapes and may trigger more
compilation or eager fallback; validate the intended frame count before
benchmarking or deployment.
The initial FP8 path targets `inference.py`; `inference_sp.py` rejects the flag
until TorchAO tensor-subclass behavior is validated with Ulysses collectives.
### NVFP4 Inference
Edit `configs/nvfp4/inference_nvfp4.yaml` to set:
- `data.data_path`: prompt folder.
- `checkpoints.generator_ckpt`: AR or base generator checkpoint.
- `checkpoints.lora_ckpt`: optional DMD LoRA checkpoint.
- `output_folder`: output video directory.
- `num_samples`: number of sampled videos per prompt.
Run:
```bash
torchrun --standalone --nnodes=1 --nproc_per_node=4 inference.py \
--config_path configs/nvfp4/inference_nvfp4.yaml
```
For single-GPU inference, use `python` directly:
```bash
python inference.py --config_path configs/nvfp4/inference_nvfp4.yaml
```
There are two recommended checkpoint styles for NVFP4 inference:
FourOverSix compact/materialized NVFP4 checkpoint:
```yaml
checkpoints:
generator_ckpt: /path/to/model_4o6.pt
merge_lora: false
model_quant: true
model_quant_use_transformer_engine: false
```
TransformerEngine runtime quantization from merged BF16 weights:
```yaml
checkpoints:
generator_ckpt: /path/to/model_te.pt
merge_lora: false
model_quant: true
model_quant_use_transformer_engine: true
```
Do not set `model_quant_use_transformer_engine: true` when loading a FourOverSix
materialized checkpoint. FourOverSix checkpoints store `quantized_weight_*`
buffers and can only be loaded by the FourOverSix path. TransformerEngine
inference should load merged BF16 weights and quantize them at runtime.
NVFP4 inference notes:
- `model_quant` enables generator NVFP4 inference. For regular BF16
checkpoints, it quantizes/materializes weights during startup; for pre-saved
FourOverSix checkpoints, the checkpoint already contains materialized weights.
- `merge_lora` merges the LoRA checkpoint into the base generator before
quantized materialization. Set it to `false` when `generator_ckpt` already
points to a merged export from `scripts/save_merged_nvfp4_generator.py`.
- `inference.kv_quant` enables FP4 KV-cache storage; the fused dequant extension
from `utils/kernel` must be built first.
- `inference.streaming_vae`, `inference.async_vae`, `inference.vae_type`, and
`inference.vae_device` control streaming or asynchronous VAE decode.
- `torch_compile` can be set to `auto`, `true`, or `false`; the default config
uses `auto` with safe error suppression.
### Sequence-parallel (SP) inference
`inference_sp.py` drives **Ulysses sequence-parallel** sampling for WAN (see `configs/inference_sp.yaml` for `sp_size`, `dp_size`, prompts, checkpoints, and the usual `inference.*` knobs). Launch one process per GPU with **`--nproc_per_node` equal to `sp_size × dp_size`** (the shipped example sets `sp_size: 4` and `dp_size: 1`, so four ranks).
```bash
torchrun --nproc_per_node=4 inference_sp.py --config_path configs/inference_sp.yaml
```
## Utilities
Inspect SP VAE halo windows:
```bash
python scripts/compute_sp_vae_chunk_halo.py --config configs/train_ar.yaml
```
Decode saved VAE latents:
```bash
python scripts/decode_vae_latents.py --help
python scripts/decode_lightvae_latents.py --help
```
example/long_example.txt
+8
View File
@@ -0,0 +1,8 @@
A close-up shot of a ceramic teacup slowly pouring water into a glass mug. The water flows smoothly from the spout of the teacup into the mug, creating gentle ripples as it fills up. Both cups have detailed textures, with the teacup having a matte finish and the glass mug showcasing clear transparency. The background is a blurred kitchen countertop, adding context without distracting from the central action. The pouring motion is fluid and natural, emphasizing the interaction between the two cups.
A dynamic and chaotic scene in a dense forest during a heavy rainstorm, capturing a real girl frantically running through the foliage. Her wild hair flows behind her as she sprints, her arms flailing and her face contorted in fear and desperation. Behind her, various animals—rabbits, deer, and birds—are also running, creating a frenzied atmosphere. The girl's clothes are soaked, clinging to her body, and she is screaming and shouting as she tries to escape. The background is a blur of greenery and rain-drenched trees, with occasional glimpses of the darkening sky. A wide-angle shot from a low angle, emphasizing the urgency and chaos of the moment.
A dynamic over-the-shoulder perspective of a chef meticulously plating a dish in a bustling kitchen. The chef, a middle-aged man with a neatly trimmed beard and focused expression, deftly arranges ingredients on a pristine white plate. His hands move with precision, each gesture deliberate and practiced. The background shows a crowded kitchen with steaming pots, whirring blenders, and the clatter of utensils. Bright lights highlight the scene, casting shadows across the busy workspace. The camera angle captures the chef's detailed work from behind, emphasizing his skill and dedication.
A dramatic and dynamic scene in the style of a disaster movie, depicting a powerful tsunami rushing through a narrow alley in Bulgaria. The water is turbulent and chaotic, with waves crashing violently against the walls and buildings on either side. The alley is lined with old, weathered houses, their facades partially submerged and splintered. The camera angle is low, capturing the full force of the tsunami as it surges forward, creating a sense of urgency and danger. People can be seen running frantically, adding to the chaos. The background features a distant horizon, hinting at the larger scale of the tsunami. A dynamic, sweeping shot from a low-angle perspective, emphasizing the movement and intensity of the event.
A playful raccoon is seen playing an electronic guitar, strumming the strings with its front paws. The raccoon has distinctive black facial markings and a bushy tail. It sits comfortably on a small stool, its body slightly tilted as it focuses intently on the instrument. The setting is a cozy, dimly lit room with vintage posters on the walls, adding a retro vibe. The raccoon's expressive eyes convey a sense of joy and concentration. Medium close-up shot, focusing on the raccoon's face and hands interacting with the guitar.
A single white sheep bending down to drink water from a calm river. The sheep has fluffy wool, long curved horns, and soft brown eyes. It is positioned near the riverbank, with its front legs partially submerged in the clear water. The river flows gently, reflecting the surrounding greenery and blue sky. The background shows lush grass and trees along the riverbank, creating a serene pastoral landscape. The sheep's body is slightly tilted as it bends down to drink, emphasizing a natural and tranquil motion. Medium close-up shot focusing on the sheep and the river.
3D animation of a small, round, fluffy creature with big, expressive eyes explores a vibrant, enchanted forest. The creature, a whimsical blend of a rabbit and a squirrel, has soft blue fur and a bushy, striped tail. It hops along a sparkling stream, its eyes wide with wonder. The forest is alive with magical elements: flowers that glow and change colors, trees with leaves in shades of purple and silver, and small floating lights that resemble fireflies. The creature stops to interact playfully with a group of tiny, fairy-like beings dancing around a mushroom ring. The creature looks up in awe at a large, glowing tree that seems to be the heart of the forest.
Animated scene features a close-up of a short fluffy monster kneeling beside a melting red candle. The art style is 3D and realistic, with a focus on lighting and texture. The mood of the painting is one of wonder and curiosity, as the monster gazes at the flame with wide eyes and open mouth. Its pose and expression convey a sense of innocence and playfulness, as if it is exploring the world around it for the first time. The use of warm colors and dramatic lighting further enhances the cozy atmosphere of the image.
fouroversix/.github/workflows/_build.yml
Vendored
+215
View File
@@ -0,0 +1,215 @@
name: ~Build wheel template
on:
workflow_call:
inputs:
runs-on:
description: "The runner to use for the build"
required: true
type: string
python-version:
description: "The Python version to use for the build"
required: true
type: string
cuda-version:
description: "The CUDA version to use for the build"
required: true
type: string
torch-version:
description: "The PyTorch version to use for the build"
required: true
type: string
cxx11_abi:
description: "The C++11 ABI to use for the build"
required: true
type: string
upload-to-release:
description: "Upload wheel to this release"
required: false
type: boolean
default: false
release-version:
description: "Upload wheel to this release"
required: false
type: string
defaults:
run:
shell: bash -x -e -u -o pipefail {0}
jobs:
build-wheel:
runs-on: ${{ inputs.runs-on }}
name: Build wheel (${{ inputs.release-version }}-${{ inputs.python-version }}-${{ inputs.cuda-version }}-${{ inputs.torch-version }}-${{ inputs.cxx11_abi }})
steps:
- name: Checkout
uses: actions/checkout@v5
with:
ref: ${{ inputs.release-version }}
submodules: recursive
- name: Set up Python
uses: actions/setup-python@v5
with:
python-version: ${{ inputs.python-version }}
- name: Set CUDA and PyTorch versions
run: |
echo "MATRIX_CUDA_VERSION=$(echo ${{ inputs.cuda-version }} | awk -F \. {'print $1 $2'})" >> $GITHUB_ENV
echo "MATRIX_TORCH_VERSION=$(echo ${{ inputs.torch-version }} | awk -F \. {'print $1 "." $2'})" >> $GITHUB_ENV
echo "WHEEL_CUDA_VERSION=$(echo ${{ inputs.cuda-version }} | awk -F \. {'print $1'})" >> $GITHUB_ENV
echo "MATRIX_PYTHON_VERSION=$(echo ${{ inputs.python-version }} | awk -F \. {'print $1 $2'})" >> $GITHUB_ENV
- name: Free up disk space
if: ${{ runner.os == 'Linux' }}
# https://github.com/easimon/maximize-build-space/blob/master/action.yml
# https://github.com/easimon/maximize-build-space/tree/test-report
run: |
sudo rm -rf /usr/share/dotnet
sudo rm -rf /usr/local/lib/android
sudo rm -rf /opt/ghc
sudo rm -rf /opt/hostedtoolcache/CodeQL
sudo docker image prune --all --force
- name: Set up swap space
if: runner.os == 'Linux'
uses: pierotofy/set-swap-space@v1.0
with:
swap-size-gb: 10
- name: Install CUDA ${{ inputs.cuda-version }}
if: ${{ inputs.cuda-version != 'cpu' }}
uses: Jimver/cuda-toolkit@v0.2.29
id: cuda-toolkit
with:
cuda: ${{ inputs.cuda-version }}
linux-local-args: '["--toolkit"]'
# default method is "local", and we're hitting some error with caching for CUDA 11.8 and 12.1
# method: ${{ (inputs.cuda-version == '11.8.0' || inputs.cuda-version == '12.1.0') && 'network' || 'local' }}
method: "network"
sub-packages: '["nvcc"]'
non-cuda-sub-packages: '["libcublas-dev", "libcusolver-dev", "libcusparse-dev"]'
- name: Install PyTorch ${{ inputs.torch-version }}+cu${{ inputs.cuda-version }}
run: |
pip install --upgrade pip
# With python 3.13 and torch 2.5.1, unless we update typing-extensions, we get error
# AttributeError: attribute '__default__' of 'typing.ParamSpec' objects is not writable
pip install typing-extensions==4.12.2
# detect if we're on ARM
if [ "$(uname -m)" = "aarch64" ] || [ "$(uname -m)" = "arm64" ]; then
PLAT=linux_aarch64
else
PLAT=manylinux_2_27_x86_64.manylinux_2_28_x86_64
fi
echo "PLAT=$PLAT" >> $GITHUB_ENV
if [[ ${{ inputs.torch-version }} == *"dev"* ]]; then
# pip install --no-cache-dir --pre torch==${{ inputs.torch-version }} --index-url https://download.pytorch.org/whl/nightly/cu${MATRIX_CUDA_VERSION}
# Can't use --no-deps because we need cudnn etc.
# Hard-coding this version of pytorch-triton for torch 2.9.0.dev20250904
pip install jinja2
TRITON_URL=https://download.pytorch.org/whl/nightly/pytorch_triton-3.4.0%2Bgitf7888497-cp${MATRIX_PYTHON_VERSION}-cp${MATRIX_PYTHON_VERSION}-${PLAT}.whl
TORCH_URL=https://download.pytorch.org/whl/nightly/cu${MATRIX_CUDA_VERSION}/torch-${{ inputs.torch-version }}%2Bcu${MATRIX_CUDA_VERSION}-cp${MATRIX_PYTHON_VERSION}-cp${MATRIX_PYTHON_VERSION}-manylinux_2_28_$(uname -m).whl
pip install --no-cache-dir --pre "${TRITON_URL}"
pip install --no-cache-dir --pre "${TORCH_URL}"
else
pip install --no-cache-dir torch==${{ inputs.torch-version }} --index-url https://download.pytorch.org/whl/cu${MATRIX_CUDA_VERSION}
fi
nvcc --version
python --version
python -c "import torch; print('PyTorch:', torch.__version__)"
python -c "import torch; print('CUDA:', torch.version.cuda)"
python -c "from torch.utils import cpp_extension; print (cpp_extension.CUDA_HOME)"
- name: Restore build cache
uses: actions/cache/restore@v4
with:
path: build.tar
key: build-${{ inputs.release-version }}-${{ inputs.python-version }}-${{ inputs.cuda-version }}-${{ inputs.torch-version }}-${{ inputs.cxx11_abi }}-${{ github.run_number }}-${{ github.run_attempt }}
restore-keys: |
build-${{ inputs.release-version }}-${{ inputs.python-version }}-${{ inputs.cuda-version }}-${{ inputs.torch-version }}-${{ inputs.cxx11_abi }}-
- name: Unpack build cache
run: |
echo ::group::Adjust timestamps
sudo find / -exec touch -t 197001010000 {} + || true
echo ::endgroup::
if [ -f build.tar ]; then
find . -mindepth 1 -maxdepth 1 ! -name 'build.tar' -exec rm -rf {} +
tar -xpvf build.tar -C .
else
echo "No build.tar found, skipping"
fi
ls -al ./
ls -al build/ || true
ls -al csrc/ || true
- name: Build wheel
id: build_wheel
run: |
pip install -U ninja packaging psutil setuptools wheel
export PATH=/usr/local/nvidia/bin:/usr/local/nvidia/lib64:$PATH
export LD_LIBRARY_PATH=/usr/local/nvidia/lib64:/usr/local/cuda/lib64:$LD_LIBRARY_PATH
# Limit MAX_JOBS otherwise the github runner goes OOM
# nvcc 11.8 can compile with 2 jobs, but nvcc 12.3 goes OOM
export MAX_JOBS=$([ "$MATRIX_CUDA_VERSION" == "129" ] && echo 1 || echo 2)
export NVCC_THREADS=2
export FORCE_BUILD=1
export FORCE_CXX11_ABI=${{ inputs.cxx11_abi == 'TRUE' && '1' || '0' }}
# 5h timeout since GH allows max 6h and we want some buffer
EXIT_CODE=0
timeout 5h python setup.py bdist_wheel --dist-dir=dist || EXIT_CODE=$?
if [ $EXIT_CODE -eq 0 ]; then
tmpname=cu${WHEEL_CUDA_VERSION}torch${MATRIX_TORCH_VERSION}cxx11abi${{ inputs.cxx11_abi }}
wheel_name=$(ls dist/*whl | xargs -n 1 basename | sed "s/-/+$tmpname-/2")
ls dist/*whl |xargs -I {} mv {} dist/${wheel_name}
echo "wheel_name=${wheel_name}" >> $GITHUB_ENV
fi
# Store exit code in GitHub env for later steps
echo "build_exit_code=$EXIT_CODE" | tee -a "$GITHUB_OUTPUT"
# Do not fail the job if timeout killed the build
exit $EXIT_CODE
- name: Log build logs after timeout
if: always() && steps.build_wheel.outputs.build_exit_code == 124
run: |
ls -al ./
tar -cvf build.tar . --atime-preserve=replace
- name: Save build cache timeout
if: always() && steps.build_wheel.outputs.build_exit_code == 124
uses: actions/cache/save@v4
with:
key: build-${{ inputs.release-version }}-${{ inputs.python-version }}-${{ inputs.cuda-version }}-${{ inputs.torch-version }}-${{ inputs.cxx11_abi }}-${{ github.run_number }}-${{ github.run_attempt }}
path: build.tar
- name: Log Built Wheels
run: |
ls dist
- name: Get Release with tag
id: get_current_release
uses: joutvhu/get-release@v1
with:
tag_name: ${{ inputs.release-version }}
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
- name: Upload Release Asset
id: upload_release_asset
if: inputs.upload-to-release
uses: actions/upload-release-asset@v1
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
with:
upload_url: ${{ steps.get_current_release.outputs.upload_url }}
asset_path: ./dist/${{env.wheel_name}}
asset_name: ${{env.wheel_name}}
asset_content_type: application/*
fouroversix/.github/workflows/build.yml
Vendored
+47
View File
@@ -0,0 +1,47 @@
name: Build wheels
on:
workflow_dispatch:
inputs:
runs-on:
description: "The runner to use for the build"
required: true
type: string
default: ubuntu-22.04
python-version:
description: "The Python version to use for the build"
required: true
type: string
cuda-version:
description: "The CUDA version to use for the build"
required: true
type: string
torch-version:
description: "The PyTorch version to use for the build"
required: true
type: string
cxx11_abi:
description: "Enable torch flag C++11 ABI (TRUE/FALSE)"
required: true
type: string
upload-to-release:
description: "Upload wheel to this release"
required: false
type: boolean
default: false
release-version:
description: "Upload wheel to this release"
required: false
type: string
jobs:
build-wheels:
uses: ./.github/workflows/_build.yml
with:
runs-on: ${{ inputs.runs-on }}
python-version: ${{ inputs.python-version }}
cuda-version: ${{ inputs.cuda-version }}
torch-version: ${{ inputs.torch-version }}
cxx11_abi: ${{ inputs.cxx11_abi }}
upload-to-release: ${{ inputs.upload-to-release }}
release-version: ${{ inputs.release-version }}
fouroversix/.github/workflows/pre-commit.yml
Vendored
+16
View File
@@ -0,0 +1,16 @@
name: Lint
on:
push:
branches:
- main
jobs:
lint:
runs-on: ubuntu-latest
strategy:
matrix:
python: ["3.13"]
steps:
- uses: actions/checkout@v4
- uses: astral-sh/ruff-action@v3
fouroversix/.github/workflows/publish.yml
Vendored
+101
View File
@@ -0,0 +1,101 @@
# This workflow will:
# - Create a new Github release
# - Build wheels for supported architectures
# - Deploy the wheels to the Github release
# - Release the static code to PyPi
# For more information see: https://help.github.com/en/actions/language-and-framework-guides/using-python-with-github-actions#publishing-to-package-registries
name: Build wheels and deploy
on:
push:
tags:
- "v*"
jobs:
setup_release:
name: Create Release
runs-on: ubuntu-latest
outputs:
release-version: ${{ steps.extract_branch.outputs.branch }}
steps:
- name: Get the tag version
id: extract_branch
run: echo ::set-output name=branch::${GITHUB_REF#refs/tags/}
shell: bash
- name: Create Release
id: create_release
uses: actions/create-release@v1
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
with:
tag_name: ${{ steps.extract_branch.outputs.branch }}
release_name: ${{ steps.extract_branch.outputs.branch }}
build_wheels:
name: Build Wheel
needs: setup_release
strategy:
fail-fast: false
matrix:
# Using ubuntu-22.04 instead of 24.04 for more compatibility (glibc). Ideally we'd use the
# manylinux docker image, but I haven't figured out how to install CUDA on manylinux.
os: [ubuntu-22.04, ubuntu-22.04-arm]
python-version: ["3.10", "3.11", "3.12", "3.13", "3.14"]
torch-version: ["2.8.0", "2.9.1", "2.10.0"]
cuda-version: ["12.9.1", "13.0.2"]
# We need separate wheels that either uses C++11 ABI (-D_GLIBCXX_USE_CXX11_ABI) or not.
# Pytorch wheels currently don't use it, but nvcr images have Pytorch compiled with C++11 ABI.
# Without this we get import error (undefined symbol: _ZN3c105ErrorC2ENS_14SourceLocationESs)
# when building without C++11 ABI and using it on nvcr images.
cxx11_abi: ["FALSE", "TRUE"]
exclude:
- os: ubuntu-22.04-arm
torch-version: "2.8.0"
- os: ubuntu-22.04-arm
torch-version: "2.7.1"
# PyTorch 2.8 only supports up to CUDA 12.9
- torch-version: "2.8.0"
cuda-version: "13.0.2"
- torch-version: "2.7.1"
cuda-version: "13.0.2"
# Python 3.14 only has pre-built wheels for PyTorch 2.9 and newer
- python-version: "3.14"
torch-version: "2.7.1"
- python-version: "3.14"
torch-version: "2.8.0"
uses: ./.github/workflows/_build.yml
with:
runs-on: ${{ matrix.os }}
python-version: ${{ matrix.python-version }}
cuda-version: ${{ matrix.cuda-version }}
torch-version: ${{ matrix.torch-version }}
cxx11_abi: ${{ matrix.cxx11_abi }}
release-version: ${{ needs.setup_release.outputs.release-version }}
upload-to-release: true
publish_package:
name: Publish package
needs: [build_wheels]
runs-on: ubuntu-latest
permissions:
id-token: write
steps:
- uses: actions/checkout@v5
with:
submodules: recursive
- uses: actions/setup-python@v6
with:
python-version: "3.10"
- name: Install dependencies
run: |
pip install -U ninja packaging setuptools wheel twine
# We don't want to download anything CUDA-related here
pip install torch --index-url https://download.pytorch.org/whl/cpu
- name: Build core package
env:
SKIP_CUDA_BUILD: "1"
run: |
python setup.py sdist --dist-dir=dist
- name: Deploy
uses: pypa/gh-action-pypi-publish@release/v1
fouroversix/.gitignore
+16
View File
@@ -0,0 +1,16 @@
__pycache__
*.py[cod]
.DS_Store
.vscode
env/
trace_*.json
*.so
exp/
*.ipynb
.cache/
ptq_logs/
dist/
*.egg-info/
site/
results.db
/ptq/
fouroversix/.gitmodules
+21
View File
@@ -0,0 +1,21 @@
[submodule "third_party/cutlass"]
path = third_party/cutlass
url = https://github.com/NVIDIA/cutlass.git
[submodule "third_party/fp-quant"]
path = third_party/fp-quant
url = https://github.com/jackcook/fp-quant-fouroversix.git
[submodule "third_party/llm-awq"]
path = third_party/llm-awq
url = https://github.com/jackcook/llm-awq-fouroversix.git
[submodule "third_party/qutlass"]
path = third_party/qutlass
url = https://github.com/IST-DASLab/qutlass.git
[submodule "third_party/fast-hadamard-transform"]
path = third_party/fast-hadamard-transform
url = https://github.com/Dao-AILab/fast-hadamard-transform.git
[submodule "third_party/spinquant"]
path = third_party/spinquant
url = https://github.com/jackcook/spinquant-fouroversix.git
[submodule "third_party/flame"]
path = third_party/flame
url = https://github.com/jackcook/flame-fouroversix.git
fouroversix/LICENSE.md
+21
View File
@@ -0,0 +1,21 @@
# MIT License
Copyright (c) 2025 Jack Cook
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
fouroversix/MANIFEST.in
+12
View File
@@ -0,0 +1,12 @@
recursive-include src *.cu
recursive-include src *.h
recursive-include src *.cpp
recursive-include src *.hpp
recursive-include src *.py
recursive-include third_party/cutlass *.cu
recursive-include third_party/cutlass *.h
recursive-include third_party/cutlass *.cuh
recursive-include third_party/cutlass *.cpp
recursive-include third_party/cutlass *.hpp
recursive-include third_party/cutlass *.inl
fouroversix/README.md
+167
View File
@@ -0,0 +1,167 @@
# Four Over Six (4/6)
[![arXiv](https://img.shields.io/badge/arXiv-2512.02010-b31b1b.svg)](https://arxiv.org/abs/2512.02010)
_Improving the accuracy of NVFP4 quantization with Adaptive Block Scaling._
![](/assets/four-over-six.png)
This repository contains kernels for efficient NVFP4 quantization and matrix multiplication, and fast post-training quantization with our method, 4/6.
If you have any questions, please get in touch or submit an issue.
## Setup
**Requirements:**
- Python version 3.10 or newer
- CUDA toolkit 12.8 or newer
- PyTorch version 2.8 or newer
**Install dependencies:**
```bash
pip install ninja packaging psutil "setuptools>=77.0.3"
```
**Install fouroversix:**
```bash
pip install fouroversix --no-build-isolation
```
Alternatively, you can compile from source:
```bash
pip install --no-build-isolation -e .
```
To speed up build times, set `CUDA_ARCHS=100` to only compile kernels for B-series GPUs (i.e. B200, GB200, GB300), or `CUDA_ARCHS=120` for RTX 50 and 60 Series GPUs (i.e. RTX 5090, RTX 6000).
Also, if you don't have a Blackwell GPU, you may use our reference implementation, which is slow but helpful for testing, by setting `SKIP_CUDA_BUILD=1` before running `pip install`.
### PTQ Experiments
To run PTQ experiments, make sure to install our test dependencies using either:
```bash
pip install "fouroversix[evals]" --no-build-isolation
# Or, if installing from source:
pip install --no-build-isolation -e ".[evals]"
```
Also, make sure all submodules are pulled and up to date:
```bash
git submodule update --init
```
Then, install dependencies for each PTQ method as needed, following the instructions [here](/docs/ptq.md).
## API
### Quantize a Model to NVFP4
```python
from fouroversix import ModelQuantizationConfig, quantize_model
from transformers import AutoModelForCausalLM
# NVFP4 using 4/6 with MSE block selection
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B")
quantize_model(model)
# Standard NVFP4 round-to-nearest quantization
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B")
config = ModelQuantizationConfig(scale_rule="static_6")
quantize_model(model, config)
```
### Quantize a Tensor to NVFP4
Check the `quantize_to_fp4` [arguments](https://github.com/mit-han-lab/fouroversix/blob/f1b78701c753ea49c091ac39d85c5753b703f5ca/src/fouroversix/frontend.py#L72) for more details about how you can enable certain features during quantization, such as stochastic rounding or 2D block quantization.
```python
import torch
from fouroversix import QuantizationConfig, quantize_to_fp4
x = torch.randn(1024, 1024, dtype=torch.bfloat16, device="cuda")
x_quantized = quantize_to_fp4(x)
# Standard NVFP4 round-to-nearest quantization
config = QuantizationConfig(scale_rule="static_6")
x_quantized = quantize_to_fp4(x, config)
```
### Multiply Two NVFP4 Tensors
```python
from fouroversix import fp4_matmul
# a and b can be either high-precision BF16 tensors, in which case they will be
# quantized, or low-precision QuantizedTensors if you've already quantized them
# yourself.
out = fp4_matmul(a, b)
```
## PTQ Evaluation with LM Evaluation Harness
```bash
# Round-to-nearest quantization with 4/6:
python -m scripts.ptq --model-name meta-llama/Llama-3.2-1B --ptq-method rtn --task wikitext
# Standard NVFP4 round-to-nearest (RTN) quantization:
python -m scripts.ptq --model-name meta-llama/Llama-3.2-1B --ptq-method rtn --task wikitext --a-scale-rule static_6 --w-scale-rule static_6
# AWQ with 4/6:
python -m scripts.ptq --model-name meta-llama/Llama-3.2-1B --ptq-method awq --task wikitext
# High-precision baseline, no NVFP4 quantization:
python -m scripts.ptq --model-name meta-llama/Llama-3.2-1B --ptq-method high_precision --task wikitext
```
If you would prefer not to worry about setting up your local environment, or about acquiring a Blackwell GPU to run your experiments faster, you may run PTQ experiments on [Modal](https://modal.com/) by adding the `--modal` flag, and optionally the `--detach` flag which will enable you to CTRL+C.
The first time you launch experiments on Modal, it may take several minutes to build everything, but following commands will reuse the cached images.
## Notes
This repository contains three implementations of NVFP4 quantization, each of which has various limitations:
- [CUDA](/src/fouroversix/csrc): Supports most but not all operations needed for efficient NVFP4 training. More operations will be added soon. Requires a Blackwell GPU.
- [Triton](/src/fouroversix/quantize/triton_kernel.py): Supports all operations needed for efficient NVFP4 training, including stochastic rounding, the random Hadamard transform, transposed inputs, and 2D block scaling. Requires a Blackwell GPU.
- [PyTorch](/src/fouroversix/quantize/reference.py): A reference implementation written in PyTorch that can run on any GPU. May have some educational value. Should not be used in real-world use cases.
When used with 4/6, these implementations have subtle numerical differences which can cause results to differ slightly, but not in a way that should cause uniformly worse performance for any of them.
For more details, see [here](https://github.com/mit-han-lab/fouroversix/blob/6bb13a8fc3b690154d11a1d6477bb6c2d09799e8/tests/test_correctness.py#L124-L132).
Our `quantize_to_fp4` function will automatically select one of these backends based on your GPU and the quantization parameters you select.
If you would like to force selection of a specific backend, you may specify it by setting `backend=QuantizeBackend.cuda` in the quantization config passed to `quantize_to_fp4`, or `quantize_backend=QuantizeBackend.cuda` in the layer and model configs passed to `quantize_model`.
## Contributing
We welcome contributions to our repository, but get in touch before making any substantial changes.
Also, please make sure any code changes are compliant with our linter:
```bash
ruff check
```
## Citation
Please use the following BibTeX entry to cite this work:
```bibtex
@misc{cook2025sixaccuratenvfp4quantization,
title={Four Over Six: More Accurate NVFP4 Quantization with Adaptive Block Scaling},
author={Jack Cook and Junxian Guo and Guangxuan Xiao and Yujun Lin and Song Han},
year={2025},
eprint={2512.02010},
archivePrefix={arXiv},
primaryClass={cs.CL},
url={https://arxiv.org/abs/2512.02010},
}
```
## License
This repository is available under the MIT license.
See the [LICENSE.md](/LICENSE.md) file for details.
fouroversix/docs/index.md
+3
View File
@@ -0,0 +1,3 @@
# Four Over Six
See [Quantization](/quantization) for details on quantizing tensors to FP4, and [Matrix Multiplication](/matmul) for performing matrix multiplication with FP4 tensors.
fouroversix/docs/matmul.md
+3
View File
@@ -0,0 +1,3 @@
# Matrix Multiplication
::: fouroversix.fp4_matmul
fouroversix/docs/ptq.md
+69
View File
@@ -0,0 +1,69 @@
# Running PTQ Experiments
## Install Dependencies
Before doing anything, make sure you've installed fouroversix with our test dependencies:
```bash
pip install -e .[evals] --no-build-isolation
```
Also, make sure you've cloned all of our submodules:
```bash
git submodule update --init
```
Then, depending on which PTQ method you would like to test, you may need to run some additional commands.
### AWQ
```bash
pip install --no-deps third_party/llm-awq
```
### GPTQ
1. Install Fast Hadamard Transform
```bash
pip install --no-build-isolation third_party/fast-hadamard-transform
```
2. Install QuTLASS
```bash
pip install --no-build-isolation third_party/qutlass
```
3. Install FP-Quant
```bash
pip install third_party/fp-quant/inference_lib
```
### High Precision
No installation necessary!
### Round-to-Nearest (RTN)
No installation necessary!
### SmoothQuant
No installation necessary!
### SpinQuant
1. Install Fast Hadamard Transform
```bash
pip install --no-build-isolation third_party/fast-hadamard-transform
```
2. Downgrade Transformers if your installation is up-to-date
```bash
pip install "transformers<5.0"
```
fouroversix/docs/quantization.md
+3
View File
@@ -0,0 +1,3 @@
# Quantization
::: fouroversix.quantize_to_fp4
fouroversix/mkdocs.yml
+8
View File
@@ -0,0 +1,8 @@
site_name: Four Over Six Documentation
theme:
name: material
plugins:
- search
- mkdocstrings
fouroversix/pyproject.toml
+110
View File
@@ -0,0 +1,110 @@
[build-system]
requires = [
"ninja",
"packaging",
"psutil",
"setuptools>=77.0.3",
]
build-backend = "setuptools.build_meta"
[project]
name = "fouroversix"
dynamic = ["version"]
authors = [
{ name="Jack Cook", email="cookj@mit.edu" },
{ name="Junxian Guo", email="junxian@mit.edu" },
]
description = "More Accurate FP4 Quantization with Adaptive Block Scaling"
readme = "README.md"
requires-python = ">=3.10"
license = "MIT"
license-files = ["LICENSE.md"]
dependencies = [
"torch>=2.7.0",
]
[project.optional-dependencies]
docs = [
"mkdocs~=1.6.1",
"mkdocs-material~=9.7.4",
"mkdocstrings[python]~=1.0.3",
]
evals = [
"inspect-ai~=0.3.186",
"inspect-evals~=0.3.106",
"lm-eval[hf]~=0.4.11",
"modal~=1.3.5",
"openai~=2.24.0",
"pytest~=8.1.1",
"ruff~=0.15.4",
"SQLAlchemy~=2.0.48",
"transformers @ git+https://github.com/huggingface/transformers.git@22c35ca",
]
[tool.pytest.ini_options]
testpaths = ["tests"]
[tool.ruff]
line-length = 88
indent-width = 4
exclude = ["build", "dist", "env", ".venv", "third_party", "scripts/ptq/tasks"]
[tool.ruff.format]
quote-style = "double"
indent-style = "space"
[tool.ruff.lint]
ignore = [
# Various dumb rules related to docstrings
"D100",
"D103", # TODO: Remove this one
"D104",
"D107",
"D202",
"D203",
"D205",
"D212",
"A001", "A002", # Allow variables to be named `input`
"FIX002", # Allow TODO comments
"FIX004", # Allow hacks
"N812", # Allow importing torch.nn.functional as F
"PLR0402", # Allow importing torch.nn as nn
"PLC0415", # Allow imports that aren't at the top of the file
"PLR0913", # Allow functions with more than 5 arguments
"RET506", # Allow elif after raise
"S603", # Allow untrusted inputs to be passed to subprocesses
"S607", # Allow commands to be run by subprocess without providing a full path
"TD003", # Don't require issue links on TODO comments
]
select = [
"ALL",
]
[tool.ruff.lint.extend-per-file-ignores]
"setup.py" = [
"T201", # Allow print statements in setup.py
]
"scripts/**/*.py" = [
"T201", # Allow print statements in CLI tools
"TID252", # Allow relative imports from parent modules
]
"src/fouroversix/**/ops.py" = [
"I001", # Allow unsorted imports, because torch needs to be imported before fouroversix._C
]
"src/fouroversix/quantize/pytorch/reference.py" = [
"S101", # Allow assert statements in reference implementation
]
"src/fouroversix/quantize/triton/kernel.py" = [
"ANN001", # Allow missing annotations in Triton kernels
"N803", # Allow uppercase arguments to Triton kernels
"N806", # Allow uppercase variables in Triton kernels
"PLR1714", "SIM109", # Allow multiple equality comparisons in Triton kernels
]
"tests/**/test_*.py" = [
"S101", # Allow assert statements in tests
"T201", # Allow print statements in tests
]
[tool.ruff.lint.isort]
known-third-party = ["fouroversix"]
fouroversix/scripts/__init__.py
View File
fouroversix/scripts/create_test_case.py
+87
View File
@@ -0,0 +1,87 @@
from .resources import Dependency, app, get_image
img = get_image(dependencies=[Dependency.transformer_engine, Dependency.fouroversix])
with img.imports():
import torch
from fouroversix import AdaptiveBlockScalingRule, QuantizeBackend, quantize_to_fp4
from fouroversix.quantize import from_blocked
@app.function(image=img, gpu="B200")
def create_test_case(
backend_a: str = "cuda",
backend_b: str = "transformer_engine",
scale_rule: str = "mse",
) -> None:
M, N = 1024, 1024 # noqa: N806
torch.set_printoptions(precision=10)
backend_a = QuantizeBackend(backend_a)
backend_b = QuantizeBackend(backend_b)
scale_rule = AdaptiveBlockScalingRule(scale_rule)
for random_seed in range(10):
torch.manual_seed(random_seed)
x = torch.randn(M, N, dtype=torch.bfloat16, device="cuda")
out_a = quantize_to_fp4(
x,
backend=backend_a,
scale_rule=scale_rule,
)
out_b = quantize_to_fp4(
x,
backend=backend_b,
scale_rule=scale_rule,
)
x_sf_a = from_blocked(out_a.scale_factors.bfloat16(), (M, N // 16))
x_sf_b = from_blocked(out_b.scale_factors.bfloat16(), (M, N // 16))
print(f"x absmax: {x.abs().max()}")
if not torch.allclose(out_a.amax, out_b.amax):
print("Backends A and B have different amax values!")
print(f"{backend_a}: {out_a.amax}")
print(f"{backend_b}: {out_b.amax}")
return
if not torch.allclose(x_sf_a.bfloat16(), x_sf_b.bfloat16()):
mismatch_prop = (x_sf_a != x_sf_b).sum() / x_sf_a.numel()
print(
"Backends A and B have different scale factors! "
f"{mismatch_prop:.2%} mismatch",
)
[i, *_], [j, *_] = torch.where(x_sf_a != x_sf_b)
print(backend_a)
print("sf", x_sf_a[i, j])
print("e2m1", out_a.e2m1_values[i, 8 * j : 8 * (j + 1)])
print(backend_b)
print("sf", x_sf_b[i, j])
print("e2m1", out_b.e2m1_values[i, 8 * j : 8 * (j + 1)])
print("original")
print("x", x[i, 16 * j : 16 * (j + 1)])
return
if not torch.allclose(out_a.e2m1_values, out_b.e2m1_values):
mismatch_prop = (
out_a.e2m1_values != out_b.e2m1_values
).sum() / out_a.e2m1_values.numel()
print(
"Backends A and B have different e2m1 values! "
f"{mismatch_prop:.2%} mismatch",
)
[i, *_], [j, *_] = torch.where(out_a.e2m1_values != out_b.e2m1_values)
print(i, j)
print("normconst", out_a.amax)
print("sf", x_sf_a[i, j // 8])
print(backend_a)
print("e2m1", out_a.e2m1_values[i, 8 * (j // 8) : 8 * (j // 8 + 1)])
print(backend_b)
print("e2m1", out_b.e2m1_values[i, 8 * (j // 8) : 8 * (j // 8 + 1)])
print("original")
print("x", x[i, 16 * (j // 8) : 16 * (j // 8 + 1)])
return
fouroversix/scripts/generate_kernels.py
+136
View File
@@ -0,0 +1,136 @@
from __future__ import annotations
import argparse
import itertools
from dataclasses import dataclass
from pathlib import Path
SRC_DTYPE_MAP = {
"fp16": "cutlass::half_t",
"bf16": "cutlass::bfloat16_t",
}
SM = [100] # Sm100 kernels support up to
IS_NVFP4 = ["false", "true"]
IS_TRANSPOSE = ["false", "true"]
IS_RHT = ["false", "true"]
def get_fp4_quant_template(
is_nvfp4: str,
is_rht: str,
is_transpose: str,
src_dtype: str,
) -> str:
if is_nvfp4 == "false":
function_str = "run_mxfp4_quant"
elif is_nvfp4 == "true":
function_str = "run_nvfp4_quant"
else:
msg = f"Invalid is_nvfp4: {is_nvfp4}"
raise ValueError(msg)
if is_rht == "true":
function_str = f"{function_str}_rht"
return f"""#include "fp4_quant_launch_template.h"
namespace fouroversix {{
template<>
void run_fp4_quant_<{src_dtype}, {is_nvfp4}, {is_rht}, {is_transpose}>(FP4_quant_params &params, cudaStream_t stream) {{
{function_str}<{src_dtype}, {is_transpose}>(params, stream);
}}
}} // namespace fouroversix""" # noqa: E501
@dataclass
class Kernel:
"""Representation for a kernel that quantizes a tensor to FP4."""
sm: int
src_dtype: str
is_nvfp4: str
is_rht: str
is_transpose: str
op: str
@property
def template(self) -> str:
"""The kernel's template content."""
template_funcs = {
"fp4_quant": get_fp4_quant_template,
}
template_func = template_funcs[self.op]
return template_func(
is_transpose=self.is_transpose,
src_dtype=SRC_DTYPE_MAP[self.src_dtype],
is_nvfp4=self.is_nvfp4,
is_rht=self.is_rht,
)
@property
def filename(self) -> str:
"""The filename for the kernel."""
fp4_format = "nvfp4" if self.is_nvfp4 == "true" else "mxfp4"
return (
f"{self.op}_{self.src_dtype}_{fp4_format}_"
f"{'rht_' if self.is_rht == 'true' else ''}"
f"{'trans_' if self.is_transpose == 'true' else ''}sm{self.sm}.cu"
)
def get_all_kernels() -> list[Kernel]:
for op in ["fp4_quant"]:
for src_dtype, is_nvfp4, is_rht, is_transpose, sm in itertools.product(
SRC_DTYPE_MAP.keys(),
IS_NVFP4,
IS_RHT,
IS_TRANSPOSE,
SM,
):
yield Kernel(
sm=sm,
src_dtype=src_dtype,
is_rht=is_rht,
is_nvfp4=is_nvfp4,
is_transpose=is_transpose,
op=op,
)
def write_kernel(kernel: Kernel, autogen_dir: Path) -> None:
prelude = """// Splitting the different transpose modes to different files to speed up compilation.
// This file is auto-generated. See "generate_kernels.py"\n""" # noqa: E501
content = prelude + kernel.template
(autogen_dir / kernel.filename).write_text(content)
def main(output_dir: str | None) -> None:
if output_dir is None:
output_dir = (
Path(__file__).parent.parent / "src" / "fouroversix" / "csrc" / "quantize"
)
else:
output_dir = Path(output_dir)
for kernel in get_all_kernels():
write_kernel(kernel, output_dir)
if __name__ == "__main__":
parser = argparse.ArgumentParser(
prog="generate_kernels",
description="Generate the flash_attention kernels template instantiations",
)
parser.add_argument(
"-o",
"--output_dir",
required=False,
help="Where to generate the kernels will default to the current directory ",
)
args = parser.parse_args()
main(args.output_dir)
fouroversix/scripts/hadamard_code_gen.py
+111
View File
@@ -0,0 +1,111 @@
from pathlib import Path
import numpy as np
# had_16 = """
# ++++++++++++++++
# +-+-+-+-+-+-+-+-
# ++--++--++--++--
# +--++--++--++--+
# ++++----++++----
# +-+--+-++-+--+-+
# ++----++++----++
# +--+-++-+--+-++-
# ++++++++--------
# +-+-+-+--+-+-+-+
# ++--++----++--++
# +--++--+-++--++-
# ++++--------++++
# +-+--+-+-+-++-+-
# ++----++--++++--
# +--+-++--++-+--+
# """
# transformerEngine style rht matrix
had_16 = """
+++-+------+-+--
+-++++-+-+-----+
++-++-++--+--+++
+---+++--+++--+-
+++--+++---++-++
+-++--+--+--+++-
++-+-+----+-+---
+------+-+++++-+
+++-+---+++-+-++
+-++++-++-+++++-
++-++-++++-++---
+---+++-+---++-+
+++--++++++--+--
+-++--+-+-++---+
++-+-+--++-+-+++
+------++-----+-
"""
header = """
/******************************************************************************
* Copyright (c) 2023, Tri Dao.
* Adapted by Junxian Guo from https://github.com/Dao-AILab/fast-hadamard-transform/blob/master/csrc/code_gen.py
* Copyright (c) 2025, FourOverSix Team.
******************************************************************************/
// This file is auto-generated. See "hadamard_code_gen.py"\n
#pragma once
"""
template = """
namespace fouroversix {{
__device__ __forceinline__ void hadamard_mult_thread_{N}(float x[{N}]) {
float out[{N}];
{code}
#pragma unroll
for (int i = 0; i < {N}; i++) { x[i] = out[i]; }
}
}} // namespace fouroversix
"""
def string_to_array(string: str) -> np.ndarray:
# Convert strings of + and - to bool arrays
string = string.strip().replace("+", "1").replace("-", "-1").split()
return np.stack(
[
np.fromstring(" ".join(string[i]), dtype=np.int32, sep=" ")
for i in range(len(string))
],
)
def array_code_gen(arr: np.ndarray) -> str:
n = arr.shape[0]
if arr.shape[0] != arr.shape[1]:
msg = f"Hadamard matrix is not square: {arr.shape}"
raise ValueError(msg)
out = [
f"out[{i}] = "
+ " ".join([f"{'+' if arr[i, j] == 1 else '-'} x[{j}]" for j in range(n)])
+ ";"
for i in range(n)
]
return template.replace("{N}", str(n)).replace("{code}", "\n ".join(out))
def main() -> None:
output_dir = (
Path(__file__).parent.parent
/ "src"
/ "fouroversix"
/ "csrc"
/ "include"
/ "hadamard_transform_te.h"
)
output_dir.write_text(header + array_code_gen(string_to_array(had_16)))
if __name__ == "__main__":
main()
fouroversix/scripts/ptq/__init__.py
View File
fouroversix/scripts/ptq/__main__.py
+105
View File
@@ -0,0 +1,105 @@
import warnings
from typing import Any
import click
import modal
from fouroversix.utils import DataType, MatmulBackend, QuantizeBackend, ScaleRule
from ..resources import app
from .coordinators import LocalEvaluationCoordinator, ModalEvaluationCoordinator
from .utils import EvaluationFramework, PTQMethod
@click.command()
@click.option(
"--activation-scale-rule",
"--a-scale-rule",
type=ScaleRule,
default=ScaleRule.mse,
)
@click.option("--detach", is_flag=True)
@click.option("--device", type=str, default="cuda")
@click.option("--dtype", type=DataType, default=DataType.nvfp4)
@click.option(
"--eval-framework",
"-f",
type=EvaluationFramework,
default=EvaluationFramework.lm_eval,
)
@click.option("--group-name", type=str, default=None)
@click.option("--limit", type=int, default=None)
@click.option("--matmul-backend", type=MatmulBackend, default=None)
@click.option("--max-length", type=int, default=None)
@click.option("--modal", is_flag=True)
@click.option("--modal-gpu", type=str)
@click.option("--model-name", "-m", type=str, multiple=True, required=True)
@click.option("--ptq-method", "-p", type=PTQMethod, multiple=True, required=True)
@click.option("--quantize-backend", type=QuantizeBackend, default=None)
@click.option("--task", "-t", type=str, multiple=True, default=["wikitext"])
@click.option("--trust-remote-code", is_flag=True)
@click.option(
"--weight-scale-rule",
"--w-scale-rule",
type=ScaleRule,
default=ScaleRule.mse,
)
@click.option("--weight-scale-2d", "--w-scale-2d", is_flag=True)
def cli(
*,
detach: bool,
group_name: str | None,
modal_gpu: str,
**kwargs: dict[str, Any],
) -> None:
activation_scale_rule = kwargs.get("activation_scale_rule")
dtype = kwargs.get("dtype")
weight_scale_rule = kwargs.get("weight_scale_rule")
model_names = kwargs.pop("model_name")
ptq_methods = kwargs.pop("ptq_method")
tasks = kwargs.pop("task")
use_modal = kwargs.pop("modal")
# Expand shortcuts
if model_names[0] == "llamaqwen":
model_names = [
"meta-llama/Llama-3.2-1B",
"meta-llama/Llama-3.1-8B",
"meta-llama/Llama-3.1-70B",
"Qwen/Qwen3-1.7B",
"Qwen/Qwen3-8B",
"Qwen/Qwen3-32B",
]
if isinstance(tasks, tuple):
tasks = list(tasks)
if dtype == DataType.mxfp4 and (
not activation_scale_rule.is_static() or not weight_scale_rule.is_static()
):
msg = (
"MXFP4 quantization only supports static scale rules. Setting "
"activation_scale_rule and weight_scale_rule to static_6..."
)
warnings.warn(msg, stacklevel=1)
kwargs["activation_scale_rule"] = ScaleRule.static_6
kwargs["weight_scale_rule"] = ScaleRule.static_6
if use_modal:
with modal.enable_output(), app.run(detach=detach):
coordinator = ModalEvaluationCoordinator(group_name_str=group_name or "")
coordinator.start.remote(
model_names,
ptq_methods,
tasks,
modal_gpu=modal_gpu,
**kwargs,
)
else:
coordinator = LocalEvaluationCoordinator(group_name)
coordinator.start(model_names, ptq_methods, tasks, **kwargs)
if __name__ == "__main__":
cli()
fouroversix/scripts/ptq/coordinators/__init__.py
+4
View File
@@ -0,0 +1,4 @@
from .local import LocalEvaluationCoordinator
from .modal import ModalEvaluationCoordinator
__all__ = ["LocalEvaluationCoordinator", "ModalEvaluationCoordinator"]
fouroversix/scripts/ptq/coordinators/base.py
+105
View File
@@ -0,0 +1,105 @@
from abc import ABC, abstractmethod
from typing import Any
from sqlalchemy import create_engine
from sqlalchemy.orm import Session, sessionmaker
from ..experiment import Base, Experiment
from ..utils import PTQMethod
class BaseEvaluationCoordinator(ABC):
"""Base class for evaluation coordinators."""
def get_session(self) -> Session:
"""Get an SQLAlchemy session for the SQLite database."""
engine = create_engine(f"sqlite:///{self.database_path.absolute().as_posix()}")
Base.metadata.create_all(engine)
return sessionmaker(bind=engine)()
def get_tasks_to_evaluate(
self,
model_name: str,
ptq_method: PTQMethod,
tasks: list[str],
) -> list[str]:
"""
Get the tasks that should be evaluated. If a group name is set, tasks will only
be evaluated if they have not yet been evaluated for this group name, model
name, PTQ method, and task.
"""
if self.group_name is None:
return tasks
session = self.get_session()
experiments = (
session.query(Experiment)
.filter(
Experiment.group_name == self.group_name,
Experiment.model_name == model_name,
Experiment.ptq_method == ptq_method.value,
Experiment.task.in_(tasks),
)
.all()
)
return [
task
for task in tasks
if task not in [experiment.task for experiment in experiments]
]
@abstractmethod
def run_calibration_tasks(
self,
model_names: list[str],
ptq_methods: list[PTQMethod],
tasks: list[str],
**kwargs: dict[str, Any],
) -> None:
"""
Run any tasks that should be used to calibrate models for a given PTQ method
and set of parameters before running evaluation.
"""
def save_results(
self,
model_name: str,
ptq_method: PTQMethod,
kwargs: dict[str, Any],
results: list[tuple[str, str, float, dict[str, Any]]],
) -> None:
"""Save the results of a PTQ experiment to the SQLite database."""
session = self.get_session()
for task, metric_name, metric_value, full_results in results:
experiment = Experiment(
group_name=self.group_name,
model_name=model_name,
task=task,
metric_name=metric_name,
metric_value=metric_value,
ptq_method=ptq_method.value,
activation_scale_rule=kwargs.get("activation_scale_rule"),
weight_scale_rule=kwargs.get("weight_scale_rule"),
smoothquant_alpha=kwargs.get("smoothquant_alpha"),
results=full_results,
)
session.add(experiment)
print(model_name, ptq_method, task)
print(full_results)
session.commit()
@abstractmethod
def start(
self,
model_names: list[str],
ptq_methods: list[PTQMethod],
tasks: list[str],
**kwargs: dict[str, Any],
) -> None:
"""Start the evaluation coordinator."""
fouroversix/scripts/ptq/coordinators/local.py
+185
View File
@@ -0,0 +1,185 @@
import itertools
import multiprocessing
from pathlib import Path
from typing import Any
import torch
from ..evaluators import get_evaluator
from ..utils import PTQMethod
from .base import BaseEvaluationCoordinator
FOUROVERSIX_ROOT_DIR = Path(__file__).parent.parent.parent.parent
class LocalEvaluationCoordinator(BaseEvaluationCoordinator):
"""Evaluation coordinator for running PTQ experiments locally."""
def __init__(self, group_name: str | None = None) -> None:
self.database_path = FOUROVERSIX_ROOT_DIR / "results.db"
self.group_name = group_name
def evaluate(
self,
model_name: str,
ptq_method: PTQMethod,
**kwargs: dict[str, Any],
) -> dict[str, Any]:
"""Evaluate a model with a given PTQ method."""
evaluator_cls = get_evaluator(ptq_method)
return evaluator_cls().evaluate(
model_name=model_name,
save_path=FOUROVERSIX_ROOT_DIR / "ptq",
**kwargs,
)
def run_calibration_tasks(
self,
model_names: list[str],
ptq_methods: list[PTQMethod],
tasks: list[str],
task_queue: multiprocessing.Queue,
result_queue: multiprocessing.Queue,
**kwargs: dict[str, Any],
) -> None:
"""
Run any tasks that should be used to calibrate models for a given PTQ method
and set of parameters before running evaluation.
"""
experiments = 0
for model_name, ptq_method in itertools.product(model_names, ptq_methods):
tasks_to_evaluate = self.get_tasks_to_evaluate(
model_name,
ptq_method,
tasks,
)
if len(tasks_to_evaluate) == 0:
continue
evaluator_cls = get_evaluator(ptq_method)
for calibration_task_kwargs in evaluator_cls.get_calibration_tasks(
model_name,
self.get_session(),
**kwargs,
):
task_queue.put(
(model_name, ptq_method, {**kwargs, **calibration_task_kwargs}),
)
experiments += 1
for _ in range(experiments):
self.save_results(*result_queue.get())
def start(
self,
model_names: list[str],
ptq_methods: list[PTQMethod],
tasks: list[str],
*,
device: str,
**kwargs: dict[str, Any],
) -> None:
"""Start the evaluation coordinator."""
multiprocessing.set_start_method("spawn", force=True)
manager = multiprocessing.Manager()
task_queue = manager.Queue()
result_queue = manager.Queue()
# Start one worker per GPU
num_workers = torch.cuda.device_count() if device == "cuda" else 1
workers = []
for gpu_id in range(num_workers):
p = multiprocessing.Process(
target=self.worker,
args=(
f"cuda:{gpu_id}" if device == "cuda" else device,
task_queue,
result_queue,
),
)
p.start()
workers.append(p)
# Run calibration tasks if necessary for each model and PTQ method
self.run_calibration_tasks(
model_names,
ptq_methods,
tasks,
task_queue,
result_queue,
**kwargs,
)
# Run evaluation tasks after models have been calibrated
experiments = 0
for model_name, ptq_method in itertools.product(model_names, ptq_methods):
tasks_to_evaluate = self.get_tasks_to_evaluate(
model_name,
ptq_method,
tasks,
)
if len(tasks_to_evaluate) == 0:
continue
evaluator_cls = get_evaluator(ptq_method)
calibrated_kwargs = evaluator_cls.get_calibrated_kwargs(
model_name,
self.get_session(),
**kwargs,
)
task_queue.put(
(
model_name,
ptq_method,
{**kwargs, "tasks": tasks_to_evaluate, **calibrated_kwargs},
),
)
experiments += 1
# Send shutdown signals (one per worker)
for _ in range(num_workers):
task_queue.put(None)
# Collect results
for _ in range(experiments):
self.save_results(*result_queue.get())
for p in workers:
p.join()
def worker(
self,
device: str,
task_queue: multiprocessing.Queue,
result_queue: multiprocessing.Queue,
) -> None:
"""Worker process for running PTQ experiments locally."""
while True:
worker_task = task_queue.get()
if worker_task is None:
break
model_name, ptq_method, kwargs = worker_task
results = self.evaluate(
model_name,
ptq_method,
**{**kwargs, "device": device},
)
result_queue.put((model_name, ptq_method, kwargs, results))
fouroversix/scripts/ptq/coordinators/modal.py
+152
View File
@@ -0,0 +1,152 @@
import itertools
from pathlib import Path
from typing import Any
import modal
from ...resources import FOUROVERSIX_CACHE_PATH, app, cache_volume, get_image
from ..evaluators import get_evaluator
from ..utils import PTQMethod
from .base import BaseEvaluationCoordinator
@app.cls(
image=get_image(),
timeout=24 * 60 * 60,
nonpreemptible=True,
volumes={FOUROVERSIX_CACHE_PATH: cache_volume},
)
class ModalEvaluationCoordinator(BaseEvaluationCoordinator):
"""Evaluation coordinator for running PTQ experiments on Modal."""
group_name_str: str = modal.parameter()
@property
def database_path(self) -> Path:
"""Path to the SQLite database where experiment results are stored."""
return FOUROVERSIX_CACHE_PATH / "results.db"
@property
def group_name(self) -> str | None:
"""
The name of the group experiments are being run in. If this is not None and an
experiment with this group name and matching parameters has already been run,
the experiment will not be run again.
"""
# Modal doesn't allow None parameters in modal.parameter()
return self.group_name_str if self.group_name_str != "" else None
def run_calibration_tasks(
self,
model_names: list[str],
ptq_methods: list[PTQMethod],
tasks: list[str],
modal_gpu: str,
**kwargs: dict[str, Any],
) -> None:
"""
Run any tasks that should be used to calibrate models for a given PTQ method
and set of parameters before running evaluation.
"""
function_calls_with_inputs = []
for model_name, ptq_method in itertools.product(model_names, ptq_methods):
tasks_to_evaluate = self.get_tasks_to_evaluate(
model_name,
ptq_method,
tasks,
)
if len(tasks_to_evaluate) == 0:
continue
evaluator_cls = get_evaluator(ptq_method).with_options(gpu=modal_gpu)
function_calls_with_inputs.extend(
[
(
model_name,
ptq_method,
{**kwargs, **calibration_task_kwargs},
evaluator_cls().evaluate_on_modal.spawn(
model_name=model_name,
save_path=FOUROVERSIX_CACHE_PATH / "ptq",
**{
**kwargs,
"tasks": tasks_to_evaluate,
**calibration_task_kwargs,
},
),
)
for calibration_task_kwargs in evaluator_cls.get_calibration_tasks(
model_name,
self.get_session(),
**kwargs,
)
],
)
results = modal.FunctionCall.gather(
*[function_call for _, _, _, function_call in function_calls_with_inputs],
)
for (model_name, ptq_method, function_call_kwargs, _), result in zip(
function_calls_with_inputs,
results,
strict=True,
):
self.save_results(model_name, ptq_method, function_call_kwargs, result)
@modal.method()
def start(
self,
model_names: list[str],
ptq_methods: list[PTQMethod],
tasks: list[str],
modal_gpu: str,
**kwargs: dict[str, Any],
) -> None:
"""Start the evaluation coordinator."""
self.run_calibration_tasks(model_names, ptq_methods, tasks, modal_gpu, **kwargs)
models_and_ptq_methods = list(itertools.product(model_names, ptq_methods))
function_calls = []
for model_name, ptq_method in models_and_ptq_methods:
tasks_to_evaluate = self.get_tasks_to_evaluate(
model_name,
ptq_method,
tasks,
)
if len(tasks_to_evaluate) == 0:
continue
evaluator_cls = get_evaluator(ptq_method).with_options(gpu=modal_gpu)
calibrated_kwargs = evaluator_cls.get_calibrated_kwargs(
model_name,
self.get_session(),
**kwargs,
)
function_calls.append(
evaluator_cls().evaluate_on_modal.spawn(
model_name=model_name,
tasks=tasks_to_evaluate,
save_path=FOUROVERSIX_CACHE_PATH / "ptq",
**{**kwargs, **calibrated_kwargs},
),
)
all_results = modal.FunctionCall.gather(*function_calls)
for (model_name, ptq_method), results in zip(
models_and_ptq_methods,
all_results,
strict=True,
):
self.save_results(model_name, ptq_method, kwargs, results)
fouroversix/scripts/ptq/evaluators/__init__.py
+34
View File
@@ -0,0 +1,34 @@
from __future__ import annotations
from typing import TYPE_CHECKING
from ..utils import PTQMethod
from .awq import AWQEvaluator
from .gptq import GPTQEvaluator
from .high_precision import HighPrecisionEvaluator
from .rtn import RTNEvaluator
from .smoothquant import SmoothQuantEvaluator
from .spinquant import SpinQuantEvaluator
if TYPE_CHECKING:
from .evaluator import PTQEvaluator
def get_evaluator(ptq_method: PTQMethod) -> type[PTQEvaluator]:
"""Get the evaluator class for the given PTQ method."""
if ptq_method == PTQMethod.awq:
return AWQEvaluator
if ptq_method == PTQMethod.gptq:
return GPTQEvaluator
if ptq_method == PTQMethod.high_precision:
return HighPrecisionEvaluator
if ptq_method == PTQMethod.rtn:
return RTNEvaluator
if ptq_method == PTQMethod.smoothquant:
return SmoothQuantEvaluator
if ptq_method == PTQMethod.spinquant:
return SpinQuantEvaluator
msg = f"Unsupported PTQ method: {ptq_method}"
raise ValueError(msg)
fouroversix/scripts/ptq/evaluators/awq.py
+134
View File
@@ -0,0 +1,134 @@
from __future__ import annotations
from pathlib import Path
from typing import TYPE_CHECKING, Any
import torch
import torch.nn as nn
import torch.nn.functional as F
from fouroversix import FourOverSixLinear, ModelQuantizationConfig
from fouroversix.model.quantize import QuantizedModule
from ...resources import (
FOUROVERSIX_CACHE_PATH,
Dependency,
app,
cache_volume,
get_image,
hf_secret,
)
from .rtn import RTNEvaluatorImpl
if TYPE_CHECKING:
from transformers import AutoModelForCausalLM
awq_img = get_image(dependencies=[Dependency.fouroversix, Dependency.awq])
class FourOverSixLinearForAWQ(FourOverSixLinear):
"""
Drop-in replacement for `FourOverSixLinear` that quantizes the weights and
activations during AWQ calibration.
"""
def __init__(self, *args: list[Any], **kwargs: dict[str, Any]) -> None:
super().__init__(*args, **kwargs)
self.config.keep_master_weights = True
self.high_precision = False
def apply_ptq(self) -> None:
"""
Override the parent method to do nothing, since we need the high-precision
weight when calibrating with AWQ.
"""
def forward(self, input: torch.Tensor) -> torch.Tensor:
"""
Forward pass that can optionally be run in high precision. This is used to
calculate the high-precision output to compare to during the auto-scale process
in AWQ calibration.
"""
return (
F.linear(input, self.weight, self.bias)
if self.high_precision
else super().forward(input)
)
@app.cls(
image=awq_img,
gpu="B200",
secrets=[hf_secret],
timeout=24 * 60 * 60,
volumes={FOUROVERSIX_CACHE_PATH.as_posix(): cache_volume},
)
class AWQEvaluator(RTNEvaluatorImpl):
"""Evaluate a model using AWQ."""
def quantize_model(
self,
model_name: str,
*,
device: str,
save_path: Path,
quantization_config: ModelQuantizationConfig,
trust_remote_code: bool,
) -> AutoModelForCausalLM:
"""Quantize a model using AWQ."""
import torch
from awq.quantize.pre_quant import apply_awq, run_awq
from fouroversix import quantize_model
from transformers import AutoModelForCausalLM, AutoTokenizer
# Replace FourOverSixLinear with FourOverSixLinearForAWQ
QuantizedModule.register(
nn.Linear,
replace_existing_modules_in_registry=True,
)(FourOverSixLinearForAWQ)
save_path = save_path / "awq" / model_name / quantization_config.__hash__()
if not save_path.exists():
model = AutoModelForCausalLM.from_pretrained(
model_name,
device_map=device,
trust_remote_code=trust_remote_code,
).eval()
quantize_model(model, quantization_config)
enc = AutoTokenizer.from_pretrained(
model_name,
use_fast=False,
trust_remote_code=trust_remote_code,
)
awq_results = run_awq(
model,
enc,
w_bit=16,
q_config={"q_group_size": -1, "zero_point": False},
n_samples=128,
seqlen=512,
calib_data="wikitext",
)
Path(save_path).parent.mkdir(parents=True, exist_ok=True)
torch.save(awq_results, save_path)
model = AutoModelForCausalLM.from_pretrained(
model_name,
device_map=device,
trust_remote_code=trust_remote_code,
)
# Apply AWQ
awq_results = torch.load(save_path, map_location="cuda")
apply_awq(model, awq_results)
# Quantize the model
quantize_model(model, quantization_config)
return model.to(device)
fouroversix/scripts/ptq/evaluators/evaluator.py
+204
View File
@@ -0,0 +1,204 @@
from __future__ import annotations
from abc import ABC, abstractmethod
from contextlib import nullcontext
from pathlib import Path
from typing import TYPE_CHECKING, Any
import modal
import torch
from fouroversix import (
DataType,
MatmulBackend,
ModelQuantizationConfig,
QuantizeBackend,
ScaleRule,
)
from ..utils import EvaluationFramework
if TYPE_CHECKING:
from sqlalchemy.orm import Session
from transformers import AutoConfig, AutoModelForCausalLM
class PTQEvaluator(ABC):
"""Base class for post-training quantization evaluators."""
@classmethod
def get_calibration_tasks(
cls,
model_name: str, # noqa: ARG003
session: Session, # noqa: ARG003
**kwargs: dict[str, Any], # noqa: ARG003
) -> list[dict[str, Any]]:
"""
Get the kwargs for tasks that should be used to calibrate the given model for
this PTQ method before running evaluation.
"""
return []
@classmethod
def get_calibrated_kwargs(
cls,
model_name: str, # noqa: ARG003
session: Session, # noqa: ARG003
**kwargs: dict[str, Any], # noqa: ARG003
) -> dict[str, Any]:
"""
Get the calibrated kwargs for the given model and scale rules. If this model
has not yet been calibrated with these scale rules, an error will be raised.
"""
return {}
@abstractmethod
def quantize_model(self, **kwargs: dict[str, Any]) -> AutoModelForCausalLM:
"""Quantize a model."""
def evaluate(
self,
model_name: str,
*,
device: str,
dtype: str,
eval_framework: EvaluationFramework,
limit: int | None,
max_length: int | None,
tasks: list[str],
trust_remote_code: bool = False,
disable_inference_mode: bool = False,
matmul_backend: MatmulBackend | None = None,
quantize_backend: QuantizeBackend | None = None,
weight_scale_2d: bool = False,
activation_scale_rule: ScaleRule | None = None,
weight_scale_rule: ScaleRule | None = None,
save_path: Path | None = None,
**kwargs: dict[str, Any],
) -> dict[str, Any]:
"""Evaluate a quantized model with lm-eval."""
inference_context = (
nullcontext() if disable_inference_mode else torch.inference_mode()
)
with inference_context:
model_config = AutoConfig.from_pretrained(model_name)
quantization_config = ModelQuantizationConfig(
activation_scale_rule=activation_scale_rule,
dtype=dtype,
matmul_backend=matmul_backend,
output_dtype=DataType(
(
str(model_config.dtype).replace("torch.", "")
if model_config.dtype is not None
else "bfloat16"
),
),
quantize_backend=quantize_backend,
weight_scale_2d=weight_scale_2d,
weight_scale_rule=weight_scale_rule,
)
model = self.quantize_model(
model_name=model_name,
device=device,
save_path=save_path,
quantization_config=quantization_config,
trust_remote_code=trust_remote_code,
**kwargs,
)
if eval_framework == EvaluationFramework.lm_eval:
from lm_eval import evaluator
from lm_eval.models.huggingface import HFLM
from lm_eval.tasks import TaskManager
full_results = evaluator.simple_evaluate(
model=HFLM(
pretrained=model,
device=device,
max_length=max_length,
),
tasks=tasks,
device=device,
limit=limit,
task_manager=TaskManager(
include_path=(
Path(__file__).parent.parent / "tasks"
).as_posix(),
),
)
results = []
for task in full_results["results"]:
result = full_results["results"][task]
if "acc_norm,none" in result:
metric_name = "acc_norm,none"
elif "acc,none" in result:
metric_name = "acc,none"
elif "word_perplexity,none" in result:
metric_name = "word_perplexity,none"
else:
metric_name = None
results.append(
(
task,
metric_name,
result.get(metric_name),
full_results["results"][task],
),
)
elif eval_framework == EvaluationFramework.inspect_ai:
import inspect_ai
from inspect_ai.model import Model
from inspect_ai.model._generate_config import GenerateConfig
from .utils import local_hf
config = GenerateConfig()
full_results = inspect_ai.eval(
tasks=tasks,
model=Model(local_hf(model_name, model, config), config, None),
limit=limit,
log_dir=(save_path / "inspect_ai_logs").as_posix(),
display="none",
)
results = []
for log in full_results:
metrics = {
k: v.value
for score in log.results.scores
for k, v in score.metrics.items()
}
metric_name = "accuracy" if "accuracy" in metrics else None
results.append(
(
log.eval.task,
metric_name,
metrics.get(metric_name),
metrics,
),
)
del model
torch.cuda.empty_cache()
return results
@modal.method()
def evaluate_on_modal(
self,
*args: list[Any],
**kwargs: dict[str, Any],
) -> dict[str, Any]:
"""Evaluate a quantized model on Modal."""
return self.evaluate(*args, **kwargs)
fouroversix/scripts/ptq/evaluators/gptq.py
+99
View File
@@ -0,0 +1,99 @@
import sys
from pathlib import Path
from typing import TYPE_CHECKING
import fouroversix
from fouroversix import ModelQuantizationConfig
from ...resources import (
FOUROVERSIX_CACHE_PATH,
Dependency,
app,
cache_volume,
get_image,
hf_secret,
)
from .evaluator import PTQEvaluator
if TYPE_CHECKING:
from transformers import AutoModelForCausalLM
CALIBRATION_DATASET = "wikitext"
gptq_img = get_image(
dependencies=[
Dependency.fast_hadamard_transform,
Dependency.qutlass,
Dependency.fp_quant,
Dependency.fouroversix,
],
)
@app.cls(
image=gptq_img,
gpu="B200",
secrets=[hf_secret],
timeout=24 * 60 * 60,
volumes={FOUROVERSIX_CACHE_PATH: cache_volume},
)
class GPTQEvaluator(PTQEvaluator):
"""Evaluate a model after quantizing it with GPTQ."""
def quantize_model(
self,
model_name: str,
*,
device: str,
save_path: Path,
quantization_config: ModelQuantizationConfig,
trust_remote_code: bool,
) -> "AutoModelForCausalLM":
"""Quantize a model with GPTQ."""
sys.path.extend(
[
(
Path(fouroversix.__file__).parent.parent.parent
/ "third_party"
/ "fp-quant"
).as_posix(),
],
)
from model_quant import main
from transformers import AutoModelForCausalLM
save_path = save_path / "gptq" / model_name / quantization_config.__hash__()
if not save_path.exists():
sys.argv = [
sys.argv[0],
"--model_name_or_path",
model_name,
"--dataset_name_or_path",
CALIBRATION_DATASET,
"--w_bits",
"4",
"--a_bits",
"4",
"--export_quantized_model",
"realquant",
"--format",
"nvfp",
"--gptq",
"--save_path",
save_path.as_posix(),
"--a_scale_rule",
quantization_config.activation_scale_rule.value,
"--w_scale_rule",
quantization_config.weight_scale_rule.value,
]
main()
return AutoModelForCausalLM.from_pretrained(
save_path,
device_map=device,
trust_remote_code=trust_remote_code,
)
fouroversix/scripts/ptq/evaluators/high_precision.py
+45
View File
@@ -0,0 +1,45 @@
from pathlib import Path
from fouroversix import ModelQuantizationConfig
from ...resources import (
FOUROVERSIX_CACHE_PATH,
app,
cache_volume,
get_image,
hf_secret,
)
from .evaluator import PTQEvaluator
hp_img = get_image()
with hp_img.imports():
from transformers import AutoModelForCausalLM
@app.cls(
image=hp_img,
gpu="B200",
secrets=[hf_secret],
timeout=24 * 60 * 60,
volumes={FOUROVERSIX_CACHE_PATH.as_posix(): cache_volume},
)
class HighPrecisionEvaluator(PTQEvaluator):
"""Evaluate a model while keeping it in high precision."""
def quantize_model(
self,
model_name: str,
*,
device: str,
save_path: Path, # noqa: ARG002
quantization_config: ModelQuantizationConfig, # noqa: ARG002
trust_remote_code: bool = False,
) -> "AutoModelForCausalLM":
"""Return a model without any quantization."""
return AutoModelForCausalLM.from_pretrained(
model_name,
device_map=device,
trust_remote_code=trust_remote_code,
)
fouroversix/scripts/ptq/evaluators/rtn.py
+106
View File
@@ -0,0 +1,106 @@
from __future__ import annotations
from typing import TYPE_CHECKING
from ...resources import (
FOUROVERSIX_CACHE_PATH,
app,
cache_volume,
get_image,
hf_secret,
)
from .evaluator import PTQEvaluator
if TYPE_CHECKING:
from pathlib import Path
from fouroversix import ModelQuantizationConfig
from transformers import AutoModelForCausalLM
rtn_img = get_image()
with rtn_img.imports():
from transformers import AutoConfig, AutoModelForCausalLM
try:
from transformers import FourOverSixConfig as HFFourOverSixConfig
except ImportError:
HFFourOverSixConfig = None
class RTNEvaluatorImpl(PTQEvaluator):
"""Evaluate a model using round-to-nearest quantization."""
def quantize_model(
self,
model_name: str,
*,
device: str,
save_path: Path,
quantization_config: ModelQuantizationConfig,
trust_remote_code: bool = False,
) -> AutoModelForCausalLM:
"""Quantize a model using round-to-nearest quantization."""
model_save_path = (
save_path / "rtn" / model_name / quantization_config.__hash__()
)
if not model_save_path.exists():
model_config = AutoConfig.from_pretrained(model_name)
hf_quantization_config = HFFourOverSixConfig(
activation_scale_rule=quantization_config.activation_scale_rule,
dtype=quantization_config.dtype,
matmul_backend=quantization_config.matmul_backend,
output_dtype=quantization_config.output_dtype,
quantize_backend=quantization_config.quantize_backend,
weight_scale_2d=quantization_config.weight_scale_2d,
weight_scale_rule=quantization_config.weight_scale_rule,
)
save_kwargs = {}
if hasattr(model_config, "quantization_config"):
hf_quantization_config.pre_quantized_model_config_type = str(
type(model_config),
)
save_kwargs["save_original_format"] = False
delattr(model_config, "quantization_config")
model = AutoModelForCausalLM.from_pretrained(
model_name,
device_map=device,
config=model_config,
quantization_config=hf_quantization_config,
trust_remote_code=trust_remote_code,
)
if hasattr(hf_quantization_config, "pre_quantized_model_config_type"):
delattr(hf_quantization_config, "pre_quantized_model_config_type")
model.save_pretrained(model_save_path, **save_kwargs)
else:
model = AutoModelForCausalLM.from_pretrained(
model_save_path,
device_map=device,
trust_remote_code=trust_remote_code,
)
# Fix for Inspect AI
model.name_or_path = model_name
return model
@app.cls(
image=rtn_img,
cpu=4,
memory=8 * 1024,
gpu="B200",
secrets=[hf_secret],
timeout=24 * 60 * 60,
volumes={FOUROVERSIX_CACHE_PATH.as_posix(): cache_volume},
)
class RTNEvaluator(RTNEvaluatorImpl):
"""Evaluate a model using round-to-nearest quantization."""
fouroversix/scripts/ptq/evaluators/smoothquant.py
+238
View File
@@ -0,0 +1,238 @@
from __future__ import annotations
from typing import TYPE_CHECKING, Any
from fouroversix import ModelQuantizationConfig, ScaleRule
from ...resources import FOUROVERSIX_CACHE_PATH, app, cache_volume, hf_secret
from ..experiment import Experiment
from ..utils import PTQMethod
from .rtn import RTNEvaluatorImpl, rtn_img
if TYPE_CHECKING:
from pathlib import Path
from sqlalchemy.orm import Session
with rtn_img.imports():
import torch
import torch.nn as nn
from fouroversix import (
FourOverSixLinear,
QuantizedModule,
fp4_matmul,
quantize_model,
)
from transformers import AutoModelForCausalLM
ALPHA_CANDIDATES = [x / 10 for x in range(11)]
WIKITEXT_TRAIN = "wikitext_train"
class FourOverSixLinearWithSmoothing(FourOverSixLinear):
"""
Drop-in replacement for `FourOverSixLinear` that implements SmoothQuant-style
scaling.
"""
def __init__(
self,
*args: list[Any],
smoothquant_alpha: float,
**kwargs: dict[str, Any],
) -> None:
super().__init__(*args, **kwargs)
self.smoothquant_alpha = smoothquant_alpha
def apply_ptq(self) -> None:
"""
Override the parent method to do nothing, since we need the high-precision
weight when doing PTQ with SmoothQuant.
"""
def forward(self, input: torch.Tensor) -> torch.Tensor:
"""Forward pass with SmoothQuant-style scaling."""
out = torch.empty(
*input.shape[:-1],
self.weight.shape[0],
device=input.device,
dtype=self.config.output_dtype.torch_dtype(),
)
fprop_activation_config = self.config.get_activation_config()
fprop_weight_config = self.config.get_weight_config(
block_scale_2d=self.config.weight_scale_2d,
)
for i in range(input.shape[0]):
s = (input[i].abs().max(dim=0).values ** self.smoothquant_alpha) / (
self.weight.abs().max(dim=0).values ** (1 - self.smoothquant_alpha)
)
out[i] = fp4_matmul(
input[i] / s[None, :],
self.weight * s[None, :],
out_dtype=self.config.output_dtype,
input_config=fprop_activation_config,
other_config=fprop_weight_config,
)
if self.bias is not None:
out = out + self.bias
return out
@app.cls(
image=rtn_img,
gpu="B200",
secrets=[hf_secret],
timeout=24 * 60 * 60,
volumes={FOUROVERSIX_CACHE_PATH.as_posix(): cache_volume},
)
class SmoothQuantEvaluator(RTNEvaluatorImpl):
"""Evaluate a model using SmoothQuant."""
@classmethod
def get_calibration_tasks(
cls,
model_name: str,
session: Session,
**kwargs: dict[str, Any],
) -> list[dict[str, Any]]:
"""
Get the kwargs for tasks that should be used to calibrate the given model for
this PTQ method before running evaluation.
"""
smoothquant_alpha = get_smoothquant_alpha(
model_name,
kwargs.get("activation_scale_rule"),
kwargs.get("weight_scale_rule"),
session,
)
calibration_experiments = get_calibration_experiments(
model_name,
kwargs.get("activation_scale_rule"),
kwargs.get("weight_scale_rule"),
session,
)
if smoothquant_alpha is None:
return [
{
"smoothquant_alpha": candidate_alpha,
"tasks": [WIKITEXT_TRAIN],
}
for candidate_alpha in ALPHA_CANDIDATES
if not any(
experiment.smoothquant_alpha == candidate_alpha
for experiment in calibration_experiments
)
]
return []
@classmethod
def get_calibrated_kwargs(
cls,
model_name: str,
session: Session,
**kwargs: dict[str, Any],
) -> dict[str, Any]:
"""
Get the calibrated kwargs for the given model and scale rules. If this model
has not yet been calibrated with these scale rules, an error will be raised.
"""
smoothquant_alpha = get_smoothquant_alpha(
model_name,
kwargs.get("activation_scale_rule"),
kwargs.get("weight_scale_rule"),
session,
)
if smoothquant_alpha is None:
msg = (
"SmoothQuant has not been calibrated for this combination of model and "
"scale rules"
)
raise ValueError(msg)
return {"smoothquant_alpha": smoothquant_alpha}
def quantize_model(
self,
model_name: str,
*,
device: str,
save_path: Path, # noqa: ARG002
smoothquant_alpha: float,
quantization_config: ModelQuantizationConfig,
trust_remote_code: bool,
) -> AutoModelForCausalLM:
"""Quantize a model using SmoothQuant."""
# Replace FourOverSixLinear with FourOverSixLinearWithSmoothing
QuantizedModule.register(
nn.Linear,
replace_existing_modules_in_registry=True,
)(FourOverSixLinearWithSmoothing)
model = AutoModelForCausalLM.from_pretrained(
model_name,
device_map=device,
trust_remote_code=trust_remote_code,
)
quantize_model(model, quantization_config, smoothquant_alpha=smoothquant_alpha)
return model
def get_calibration_experiments(
model_name: str,
activation_scale_rule: ScaleRule,
weight_scale_rule: ScaleRule,
db_session: Session,
) -> list[Experiment]:
return (
db_session.query(Experiment)
.filter(
Experiment.ptq_method == PTQMethod.smoothquant.value,
Experiment.task == WIKITEXT_TRAIN,
Experiment.model_name == model_name,
Experiment.activation_scale_rule == activation_scale_rule.value,
Experiment.weight_scale_rule == weight_scale_rule.value,
Experiment.smoothquant_alpha.isnot(None),
)
.all()
)
def get_smoothquant_alpha(
model_name: str,
activation_scale_rule: ScaleRule,
weight_scale_rule: ScaleRule,
session: Session,
) -> float | None:
calibration_experiments = get_calibration_experiments(
model_name,
activation_scale_rule,
weight_scale_rule,
session,
)
if not all(
any(
experiment.smoothquant_alpha == alpha
for experiment in calibration_experiments
)
for alpha in ALPHA_CANDIDATES
):
return None
return min(calibration_experiments, key=lambda x: x.metric_value).smoothquant_alpha
fouroversix/scripts/ptq/evaluators/spinquant.py
+240
View File
@@ -0,0 +1,240 @@
from __future__ import annotations
import subprocess
import sys
from pathlib import Path
from typing import TYPE_CHECKING, Any
import fouroversix
import modal
from fouroversix import ModelQuantizationConfig
from ...resources import (
FOUROVERSIX_CACHE_PATH,
Dependency,
app,
cache_volume,
get_image,
hf_secret,
)
from ..utils import get_model_size
from .evaluator import PTQEvaluator
if TYPE_CHECKING:
from transformers import AutoModelForCausalLM
spinquant_img = get_image(
dependencies=[
Dependency.fast_hadamard_transform,
Dependency.fouroversix,
Dependency.spinquant,
],
extra_pip_dependencies=["transformers<5.0"],
)
MIN_MODEL_SIZE_FOR_8xB200 = 32
SPINQUANT_STEPS = 100
SPINQUANT_ARGS = [
"--model_max_length",
"8192",
"--fp16",
"False",
"--bf16",
"True",
"--w_bits",
"4",
"--a_bits",
"4",
"--k_bits",
"16",
"--v_bits",
"16",
]
@app.cls(
image=spinquant_img,
timeout=24 * 60 * 60,
secrets=[hf_secret],
volumes={FOUROVERSIX_CACHE_PATH.as_posix(): cache_volume},
)
class SpinQuantOptimizer:
"""Optimize a model with SpinQuant."""
def optimize(
self,
model_name: str,
*,
quantization_config: ModelQuantizationConfig,
spinquant_save_path: str,
spinquant_steps: int,
) -> None:
"""Optimize a model with SpinQuant."""
subprocess.run(
[
"torchrun",
"--nnodes=1",
"--nproc_per_node=auto",
(
Path(fouroversix.__file__).parent.parent.parent
/ "third_party"
/ "spinquant"
/ "optimize_rotation.py"
).as_posix(),
"--input_model",
model_name,
"--output_dir",
spinquant_save_path,
"--output_rotation_path",
spinquant_save_path,
"--log_on_each_node",
"False",
"--per_device_train_batch_size",
"1",
"--logging_steps",
"1",
"--learning_rate",
"1.5",
"--weight_decay",
"0.",
"--lr_scheduler_type",
"cosine",
"--gradient_checkpointing",
"True",
"--save_safetensors",
"False",
"--max_steps",
str(spinquant_steps),
"--activation_scale_rule",
quantization_config.activation_scale_rule.value,
"--weight_scale_rule",
quantization_config.weight_scale_rule.value,
*SPINQUANT_ARGS,
],
check=True,
)
cache_volume.commit()
@modal.method()
def optimize_on_modal(
self,
*args: list[Any],
**kwargs: dict[str, Any],
) -> None:
"""Optimize a model with SpinQuant on Modal."""
return self.optimize(*args, **kwargs)
@app.cls(
image=spinquant_img,
timeout=24 * 60 * 60,
secrets=[hf_secret],
gpu="B200",
volumes={FOUROVERSIX_CACHE_PATH.as_posix(): cache_volume},
)
class SpinQuantEvaluator(PTQEvaluator):
"""Evaluate a quantized model with SpinQuant."""
def quantize_model(
self,
model_name: str,
*,
device: str,
save_path: Path,
quantization_config: ModelQuantizationConfig,
trust_remote_code: bool,
) -> AutoModelForCausalLM:
"""Export a quantized model with SpinQuant."""
import fouroversix
sys.path.append(
(
Path(fouroversix.__file__).parent.parent.parent
/ "third_party"
/ "spinquant"
).as_posix(),
)
from eval_utils.main import ptq_model
from transformers import AutoConfig, AutoModelForCausalLM
from utils.process_args import process_args_ptq
save_path = (
save_path
/ "spinquant"
/ (
f"{model_name}-{quantization_config.activation_scale_rule.value}"
f"-{quantization_config.weight_scale_rule.value}"
)
)
if not (save_path / "R.bin").exists():
model_is_large = get_model_size(model_name) >= MIN_MODEL_SIZE_FOR_8xB200
if model_is_large:
msg = (
"Automatic SpinQuant optimization is not supported for large "
"models. Please optimize the model manually."
)
raise RuntimeError(msg)
SpinQuantOptimizer().optimize(
model_name,
quantization_config=quantization_config,
spinquant_save_path=save_path.as_posix(),
spinquant_steps=SPINQUANT_STEPS,
)
sys.argv = [
sys.argv[0],
"--input_model",
model_name,
"--do_train",
"False",
"--do_eval",
"True",
"--per_device_eval_batch_size",
"4",
"--rotate",
"--optimized_rotation_path",
(save_path / "R.bin").as_posix(),
"--activation_scale_rule",
quantization_config.activation_scale_rule.value,
"--weight_scale_rule",
quantization_config.weight_scale_rule.value,
*SPINQUANT_ARGS,
]
config = AutoConfig.from_pretrained(model_name)
# Llama v3.2 specific: Spinquant is not compatiable with tie_word_embeddings,
# clone lm_head from embed_tokens
process_word_embeddings = False
if config.tie_word_embeddings:
config.tie_word_embeddings = False
process_word_embeddings = True
model = AutoModelForCausalLM.from_pretrained(
model_name,
config=config,
device_map=device,
trust_remote_code=trust_remote_code,
)
if process_word_embeddings:
model.lm_head.weight.data = model.model.embed_tokens.weight.data.clone()
model.to(device)
model_args, _, ptq_args = process_args_ptq()
cache_volume.reload()
model = ptq_model(ptq_args, model, model_args)
model.to(device)
return model
fouroversix/scripts/ptq/evaluators/utils.py
+100
View File
@@ -0,0 +1,100 @@
from typing import Any
import torch
from inspect_ai.model import modelapi
from inspect_ai.model._generate_config import GenerateConfig
from inspect_ai.model._providers.hf import HuggingFaceAPI
from transformers import AutoModelForCausalLM, AutoTokenizer
def set_random_seeds(seed: int | None = None) -> None:
import os
import numpy as np
from transformers import set_seed
if seed is None:
seed = np.random.default_rng().integers(2**32 - 1)
# python hash seed
os.environ["PYTHONHASHSEED"] = str(seed)
# transformers seed
set_seed(seed)
class LocalHuggingFaceAPI(HuggingFaceAPI):
"""
Wrapper around HuggingFaceAPI that allows for quantized models to be used during
evaluation.
"""
def __init__( # noqa: C901
self,
model_name: str,
model: AutoModelForCausalLM,
config: GenerateConfig | None = None,
**model_args: dict[str, Any],
) -> None:
self.model_name = model_name
self.base_url = None
self.api_key = None
self.api_key_vars = ["HF_TOKEN"]
self._apply_api_key_overrides()
if config is None:
config = GenerateConfig()
# set random seeds
if config.seed is not None:
set_random_seeds(config.seed)
# collect known model_args (then delete them so we can pass the rest on)
def collect_model_arg(name: str) -> Any | None: # noqa: ANN401
nonlocal model_args
value = model_args.get(name)
if value is not None:
model_args.pop(name)
return value
device = collect_model_arg("device")
tokenizer = collect_model_arg("tokenizer")
model_path = collect_model_arg("model_path")
tokenizer_path = collect_model_arg("tokenizer_path")
self.batch_size = collect_model_arg("batch_size")
self.chat_template = collect_model_arg("chat_template")
self.tokenizer_call_args = collect_model_arg("tokenizer_call_args")
self.enable_thinking = collect_model_arg("enable_thinking")
if self.tokenizer_call_args is None:
self.tokenizer_call_args = {}
self.hidden_states = collect_model_arg("hidden_states")
# device
if device:
self.device = device
elif torch.backends.mps.is_available():
self.device = "mps"
elif torch.cuda.is_available():
self.device = "cuda:0"
else:
self.device = "cpu"
# model
self.model = model
# tokenizer
if tokenizer:
self.tokenizer = AutoTokenizer.from_pretrained(tokenizer) # type: ignore[no-untyped-call]
elif model_path:
if tokenizer_path:
self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path) # type: ignore[no-untyped-call]
else:
self.tokenizer = AutoTokenizer.from_pretrained(model_path) # type: ignore[no-untyped-call]
else:
self.tokenizer = AutoTokenizer.from_pretrained(model_name) # type: ignore[no-untyped-call]
# LLMs generally don't have a pad token and we need one for batching
self.tokenizer.pad_token = self.tokenizer.eos_token
self.tokenizer.padding_side = "left"
@modelapi(name="local_hf")
def local_hf() -> type[LocalHuggingFaceAPI]:
return LocalHuggingFaceAPI
fouroversix/scripts/ptq/experiment.py
+22
View File
@@ -0,0 +1,22 @@
from sqlalchemy import JSON, Column, Float, Integer, String
from sqlalchemy.ext.declarative import declarative_base
Base = declarative_base()
class Experiment(Base):
"""A PTQ experiment with results."""
__tablename__ = "experiments"
id = Column(Integer, primary_key=True, autoincrement=True)
group_name = Column(String)
model_name = Column(String, nullable=False)
task = Column(String, nullable=False)
metric_name = Column(String, nullable=False)
metric_value = Column(Float, nullable=False)
ptq_method = Column(String, nullable=False)
activation_scale_rule = Column(String, nullable=False)
weight_scale_rule = Column(String, nullable=False)
smoothquant_alpha = Column(Float, nullable=True)
results = Column(JSON, nullable=False)
fouroversix/scripts/ptq/tasks/__init__.py
View File
fouroversix/scripts/ptq/tasks/wikitext_train/preprocess_wikitext.py
+48
View File
@@ -0,0 +1,48 @@
import re
def wikitext_detokenizer(doc):
string = doc["page"]
# contractions
string = string.replace("s '", "s'")
string = re.sub(r"/' [0-9]/", r"/'[0-9]/", string)
# number separators
string = string.replace(" @-@ ", "-")
string = string.replace(" @,@ ", ",")
string = string.replace(" @.@ ", ".")
# punctuation
string = string.replace(" : ", ": ")
string = string.replace(" ; ", "; ")
string = string.replace(" . ", ". ")
string = string.replace(" ! ", "! ")
string = string.replace(" ? ", "? ")
string = string.replace(" , ", ", ")
# double brackets
string = re.sub(r"\(\s*([^\)]*?)\s*\)", r"(\1)", string)
string = re.sub(r"\[\s*([^\]]*?)\s*\]", r"[\1]", string)
string = re.sub(r"{\s*([^}]*?)\s*}", r"{\1}", string)
string = re.sub(r"\"\s*([^\"]*?)\s*\"", r'"\1"', string)
string = re.sub(r"'\s*([^']*?)\s*'", r"'\1'", string)
# miscellaneous
string = string.replace("= = = =", "====")
string = string.replace("= = =", "===")
string = string.replace("= =", "==")
string = string.replace(" " + chr(176) + " ", chr(176))
string = string.replace(" \n", "\n")
string = string.replace("\n ", "\n")
string = string.replace(" N ", " 1 ")
string = string.replace(" 's", "'s")
return string
def process_results(doc, results):
(loglikelihood,) = results
# IMPORTANT: wikitext counts number of words in *original doc before detokenization*
_words = len(re.split(r"\s+", doc["page"]))
_bytes = len(doc["page"].encode("utf-8"))
return {
"word_perplexity": (loglikelihood, _words),
"byte_perplexity": (loglikelihood, _bytes),
"bits_per_byte": (loglikelihood, _bytes),
}
fouroversix/scripts/ptq/tasks/wikitext_train/wikitext_train.yaml
+16
View File
@@ -0,0 +1,16 @@
task: wikitext_train
dataset_path: EleutherAI/wikitext_document_level
dataset_name: wikitext-2-raw-v1
output_type: loglikelihood_rolling
test_split: train
doc_to_text: ""
doc_to_target: !function preprocess_wikitext.wikitext_detokenizer
process_results: !function preprocess_wikitext.process_results
should_decontaminate: true
doc_to_decontamination_query: "{{page}}"
metric_list:
- metric: word_perplexity
- metric: byte_perplexity
- metric: bits_per_byte
metadata:
version: 2.0
fouroversix/scripts/ptq/utils.py
+25
View File
@@ -0,0 +1,25 @@
from __future__ import annotations
from enum import Enum
class EvaluationFramework(str, Enum):
"""Frameworks to use when evaluating models."""
inspect_ai = "inspect_ai"
lm_eval = "lm_eval"
class PTQMethod(str, Enum):
"""Methods of post-training quantization."""
awq = "awq"
high_precision = "high_precision"
gptq = "gptq"
rtn = "rtn"
smoothquant = "smoothquant"
spinquant = "spinquant"
def get_model_size(model_name: str | None) -> float:
return float(model_name.split("-")[-1][:-1]) if model_name else 0.0
fouroversix/scripts/resources.py
+388
View File
@@ -0,0 +1,388 @@
from __future__ import annotations
import configparser
import os
import shutil
import subprocess
from enum import Enum
from pathlib import Path
from typing import TYPE_CHECKING
import modal
import tomllib
if TYPE_CHECKING:
from collections.abc import Callable
FOUROVERSIX_CACHE_PATH = Path("/fouroversix")
FOUROVERSIX_INSTALL_PATH = Path("/root/fouroversix")
KERNEL_DEV_MODE = os.getenv("KERNEL_DEV_MODE", "0") == "1"
app = modal.App("fouroversix")
cache_volume = modal.Volume.from_name("fouroversix", create_if_missing=True)
hf_secret = modal.Secret.from_name("huggingface-secret")
wandb_secret = modal.Secret.from_name("wandb-secret")
class Dependency(str, Enum):
"""Dependencies to add to the base image."""
awq = "awq"
fast_hadamard_transform = "fast_hadamard_transform"
flame = "flame"
flash_attention = "flash_attention"
fouroversix = "fouroversix"
fp_quant = "fp_quant"
qutlass = "qutlass"
spinquant = "spinquant"
transformer_engine = "transformer_engine"
class Submodule(str, Enum):
"""Submodules of Four Over Six to add to the base image."""
cutlass = "cutlass"
flame = "flame"
fast_hadamard_transform = "fast_hadamard_transform"
fp_quant = "fp_quant"
llm_awq = "llm_awq"
qutlass = "qutlass"
spinquant = "spinquant"
def has_untracked_or_unstaged_changes(self) -> bool:
"""Check if the submodule has untracked or unstaged changes."""
git_status = subprocess.run(
[
"git",
"-C",
self.get_local_path(),
"status",
"--porcelain",
],
check=False,
stderr=subprocess.DEVNULL,
stdout=subprocess.PIPE,
text=True,
)
return bool(git_status.stdout.strip())
def get_install_path(self) -> str:
"""Get the path where this submodule will be installed in the Modal image."""
return f"{FOUROVERSIX_INSTALL_PATH}/{self.get_local_path()}"
def get_local_path(self) -> str:
"""Get the path of the submodule relative to the root directory."""
return f"third_party/{self.value.replace('_', '-')}"
def get_remote_url(self) -> str:
"""Get the remote URL of the submodule."""
gitmodules_path = Path(__file__).parent.parent / ".gitmodules"
if not gitmodules_path.exists():
gitmodules_path = FOUROVERSIX_INSTALL_PATH / ".gitmodules"
with gitmodules_path.open() as f:
# Remove leading whitespace to make it a valid INI file
gitmodules_contents = "\n".join(line.lstrip() for line in f.readlines())
config = configparser.ConfigParser()
config.read_string(gitmodules_contents)
for section in config.sections():
if config[section]["path"] == self.get_local_path():
url = config[section]["url"]
break
if url.startswith("https://"):
return url
msg = f"Unsupported remote URL format: {url}"
raise ValueError(msg)
cuda_version_to_image_tag = {
"12.8": "nvcr.io/nvidia/cuda-dl-base:25.03-cuda12.8-devel-ubuntu24.04",
"12.9": "nvcr.io/nvidia/cuda-dl-base:25.06-cuda12.9-devel-ubuntu24.04",
"13.0": "nvcr.io/nvidia/cuda-dl-base:25.09-cuda13.0-devel-ubuntu24.04",
"13.1": "nvcr.io/nvidia/cuda-dl-base:25.12-cuda13.1-devel-ubuntu24.04",
}
def add_submodule(img: modal.Image, submodule: Submodule) -> modal.Image:
if submodule.has_untracked_or_unstaged_changes():
# Submodule has uncommitted changes, build image with local copy
return img.add_local_dir(
submodule.get_local_path(),
submodule.get_install_path(),
copy=True,
)
# Submodule has no uncommitted changes, download from remote to save time
return img.run_commands(
f"git clone {submodule.get_remote_url()} {submodule.get_install_path()}",
)
def install_flash_attn() -> None:
subprocess.run(
["pip", "install", "flash-attn", "--no-build-isolation"],
check=False,
)
def install_fouroversix() -> None:
subprocess.run(
[
"pip",
"install",
"--no-deps",
"--no-build-isolation",
"-e",
FOUROVERSIX_INSTALL_PATH.as_posix(),
],
check=False,
)
def install_fouroversix_non_editable() -> None:
shutil.copytree(
FOUROVERSIX_CACHE_PATH / "build",
FOUROVERSIX_INSTALL_PATH / "build",
)
subprocess.run(
["python", "setup.py", "build_ext", "--inplace"],
check=False,
)
shutil.copytree(
FOUROVERSIX_INSTALL_PATH / "build",
FOUROVERSIX_CACHE_PATH / "build",
dirs_exist_ok=True,
)
subprocess.run(
[
"pip",
"install",
"--no-deps",
"--no-build-isolation",
FOUROVERSIX_INSTALL_PATH.as_posix(),
],
check=False,
)
def install_qutlass() -> None:
subprocess.run(
[
"pip",
"install",
"--no-build-isolation",
Submodule.qutlass.get_install_path(),
],
check=False,
)
def get_image( # noqa: C901, PLR0912
dependencies: list[Dependency] | None = None,
*,
cuda_version: str = "12.9",
deploy: bool = False,
extra_env: dict[str, str] | None = None,
extra_pip_dependencies: list[str] | None = None,
include_tests: bool = False,
python_version: str = "3.13",
pytorch_version: str = "2.10.0",
run_before_copy: Callable[[modal.Image], modal.Image] | None = None,
) -> modal.Image:
if dependencies is None:
dependencies = [Dependency.fouroversix]
pyproject_path = Path(__file__).parent.parent / "pyproject.toml"
if not pyproject_path.exists():
pyproject_path = Path(__file__).parent.parent / "fouroversix" / "pyproject.toml"
with pyproject_path.open("rb") as f:
pyproject_data = tomllib.load(f)
img = (
modal.Image.from_registry(
cuda_version_to_image_tag[cuda_version],
add_python=python_version,
)
.entrypoint([])
.apt_install("clang", "git")
.uv_pip_install(*pyproject_data["build-system"]["requires"], "numpy")
.uv_pip_install(
f"torch=={pytorch_version}",
extra_index_url=(
f"https://download.pytorch.org/whl/cu{cuda_version.replace('.', '')}"
),
)
)
for dependency in dependencies:
if dependency == Dependency.awq:
img = add_submodule(img, Submodule.llm_awq).run_commands(
f"pip install --no-deps -e {Submodule.llm_awq.get_install_path()}",
)
if dependency == Dependency.fast_hadamard_transform:
img = add_submodule(img, Submodule.fast_hadamard_transform).run_commands(
f"pip install {Submodule.fast_hadamard_transform.get_install_path()} "
"--no-build-isolation",
)
if dependency == Dependency.flame:
img = (
img.apt_install("pciutils")
.uv_pip_install(
"flash-linear-attention",
"ninja",
"psutil",
"git+https://github.com/pytorch/torchtitan.git@0b44d4c",
"tyro",
"wheel",
)
.run_commands(
"git clone https://github.com/fla-org/flame.git "
f"{FOUROVERSIX_INSTALL_PATH}/third_party/flame",
f"pip install -e {FOUROVERSIX_INSTALL_PATH}/third_party/flame",
)
)
if dependency == Dependency.flash_attention:
img = img.run_function(
install_flash_attn,
cpu=64,
memory=128 * 1024,
gpu="B200",
)
if dependency == Dependency.fouroversix:
img = (
add_submodule(
img.env(
{"CUDA_ARCHS": "100", "FORCE_BUILD": "1", "MAX_JOBS": "32"},
),
Submodule.cutlass,
)
.add_local_file(
"pyproject.toml",
f"{FOUROVERSIX_INSTALL_PATH}/pyproject.toml",
copy=True,
)
.uv_pip_install(
*pyproject_data["project"]["optional-dependencies"]["evals"],
)
.add_local_file(
"setup.py",
f"{FOUROVERSIX_INSTALL_PATH}/setup.py",
copy=True,
)
.add_local_file(
"src/fouroversix/__init__.py",
f"{FOUROVERSIX_INSTALL_PATH}/src/fouroversix/__init__.py",
copy=True,
)
)
if KERNEL_DEV_MODE:
img = (
img.add_local_file(
"README.md",
f"{FOUROVERSIX_INSTALL_PATH}/README.md",
copy=True,
)
.add_local_file(
"LICENSE.md",
f"{FOUROVERSIX_INSTALL_PATH}/LICENSE.md",
copy=True,
)
.workdir(FOUROVERSIX_INSTALL_PATH)
)
img = img.add_local_dir(
"src/fouroversix/csrc",
f"{FOUROVERSIX_INSTALL_PATH}/src/fouroversix/csrc",
copy=True,
)
if not KERNEL_DEV_MODE:
img = img.run_function(install_fouroversix, cpu=32, memory=64 * 1024)
if dependency == Dependency.fp_quant:
img = add_submodule(img, Submodule.fp_quant).run_commands(
f"pip install {Submodule.fp_quant.get_install_path()}/inference_lib",
)
if dependency == Dependency.qutlass:
img = (
add_submodule(img.apt_install("cmake"), Submodule.qutlass)
.env({"MAX_JOBS": "32"})
.run_function(install_qutlass, gpu="B200", cpu=32, memory=64 * 1024)
)
if dependency == Dependency.spinquant:
img = add_submodule(img, Submodule.spinquant)
if dependency == Dependency.transformer_engine:
img = img.uv_pip_install(
"transformer_engine[pytorch]",
extra_options="--no-build-isolation",
)
if extra_pip_dependencies is not None:
img = img.uv_pip_install(*extra_pip_dependencies)
img = img.env({"HF_HOME": FOUROVERSIX_CACHE_PATH.as_posix(), **(extra_env or {})})
if run_before_copy is not None:
img = run_before_copy(img)
# Add source files after all dependencies are added so we can avoid rebuilding when
# they change
for dependency in dependencies:
if dependency == Dependency.flame:
img = (
img.add_local_dir(
"third_party/flame/custom_models",
f"{FOUROVERSIX_INSTALL_PATH}/third_party/flame/custom_models",
)
.add_local_dir(
"scripts/train/configs",
f"{FOUROVERSIX_INSTALL_PATH}/scripts/train/configs",
)
.add_local_file(
"third_party/flame/train.sh",
f"{FOUROVERSIX_INSTALL_PATH}/third_party/flame/train.sh",
)
)
if dependency == Dependency.fouroversix:
img = img.add_local_dir(
"src",
f"{FOUROVERSIX_INSTALL_PATH}/src",
copy=deploy or KERNEL_DEV_MODE,
ignore=lambda p: p.suffix == ".so",
).add_local_file(
".gitmodules",
f"{FOUROVERSIX_INSTALL_PATH}/.gitmodules",
copy=deploy or KERNEL_DEV_MODE,
)
if KERNEL_DEV_MODE:
img = img.run_function(
install_fouroversix_non_editable,
cpu=32,
memory=64 * 1024,
volumes={FOUROVERSIX_CACHE_PATH.as_posix(): cache_volume},
)
if include_tests:
img = img.add_local_dir("tests", f"{FOUROVERSIX_INSTALL_PATH}/tests")
return img
fouroversix/scripts/speedtest/__init__.py
View File
fouroversix/scripts/speedtest/matmul.py
+93
View File
@@ -0,0 +1,93 @@
from __future__ import annotations
from typing import Any
import click
import modal
from ..resources import app, get_image
img = get_image()
with img.imports():
import torch
import torch.utils.benchmark as benchmark
from fouroversix import DataType, MatmulBackend, QuantizationConfig, quantize_to_fp4
from fouroversix.matmul.frontend import AVAILABLE_BACKENDS
def run_speedtest(
*,
dtype: DataType = DataType.nvfp4,
m: int = 1024,
n: int = 1024,
k: int = 1024,
repeats: int = 100,
) -> None:
"""Test speed on a B200 on Modal."""
x = torch.randn(m, k, dtype=torch.bfloat16, device="cuda")
y = torch.randn(k, n, dtype=torch.bfloat16, device="cuda")
config = QuantizationConfig(dtype=dtype)
x_quantized = quantize_to_fp4(x, config)
y_quantized = quantize_to_fp4(y, config)
out_dtype = DataType.bfloat16
print(f"Testing with {m}x{k} @ {k}x{n}")
for backend in [MatmulBackend.cutlass, MatmulBackend.pytorch]:
backend_cls = AVAILABLE_BACKENDS[backend]
print(f"{backend.value}: ", end="")
if not backend_cls.is_available():
print("Not available")
continue
if not backend_cls.is_supported(
x_quantized,
y_quantized,
out_dtype=out_dtype,
):
print("Not supported")
continue
t = benchmark.Timer(
setup="from fouroversix import fp4_matmul",
stmt=(
"fp4_matmul(x_quantized, y_quantized, backend=backend, "
"out_dtype=out_dtype)"
),
globals={
"x_quantized": x_quantized,
"y_quantized": y_quantized,
"backend": backend,
"out_dtype": out_dtype,
},
)
print(f"{t.timeit(repeats).mean * 1000:.4f}ms")
@app.function(image=img, cpu=4, memory=8 * 1024, gpu="B200")
def run_speedtest_on_modal(**kwargs: dict[str, Any]) -> None:
run_speedtest(**kwargs)
@click.command()
@click.option("--dtype", type=DataType, default=DataType.nvfp4)
@click.option("--m", type=int, default=1024)
@click.option("--modal", is_flag=True)
@click.option("--n", type=int, default=1024)
@click.option("--k", type=int, default=1024)
@click.option("--repeats", type=int, default=100)
def cli(**kwargs: dict[str, Any]) -> None:
if kwargs.pop("modal"):
with modal.enable_output(), app.run():
run_speedtest_on_modal.remote(**kwargs)
else:
run_speedtest(**kwargs)
if __name__ == "__main__":
cli()
fouroversix/scripts/speedtest/quantize.py
+108
View File
@@ -0,0 +1,108 @@
from __future__ import annotations
from typing import Any
import click
import modal
from ..resources import Dependency, app, get_image
img = get_image(dependencies=[Dependency.transformer_engine, Dependency.fouroversix])
with img.imports():
import torch
import torch.utils.benchmark as benchmark
from fouroversix import QuantizationConfig, QuantizeBackend, RoundStyle, ScaleRule
from fouroversix.quantize.frontend import AVAILABLE_BACKENDS
def run_speedtest(
*,
block_scale_2d: bool = False,
input_shape: str = "1024,1024",
repeats: int = 100,
rht: bool = False,
round_style: str = "nearest",
scale_rule: str = "mse",
transpose: bool = False,
) -> None:
"""Test speed on a B200 on Modal."""
input_shape = tuple(int(dim.strip()) for dim in input_shape.split(","))
x = torch.randn(input_shape, dtype=torch.bfloat16, device="cuda")
print("Testing with config:")
print(f"- block_scale_2d: {block_scale_2d}")
print(f"- input_shape: {input_shape}")
print(f"- rht: {rht}")
print(f"- round_style: {round_style}")
print(f"- scale_rule: {scale_rule}")
print(f"- transpose: {transpose}")
print()
for backend in [
QuantizeBackend.cuda,
QuantizeBackend.transformer_engine,
QuantizeBackend.triton,
QuantizeBackend.pytorch,
]:
config = QuantizationConfig(
backend=backend,
block_scale_2d=block_scale_2d,
rht=rht,
round_style=RoundStyle(round_style),
scale_rule=ScaleRule(scale_rule),
transpose=transpose,
)
backend_cls = AVAILABLE_BACKENDS[backend]
print(f"{backend.value}: ", end="")
if not backend_cls.is_available():
print("Not available")
continue
if not backend_cls.is_supported(x, config):
print("Not supported")
continue
config = QuantizationConfig(
backend=backend,
rht=rht,
round_style=RoundStyle(round_style),
scale_rule=ScaleRule(scale_rule),
)
t = benchmark.Timer(
setup="from fouroversix import quantize_to_fp4",
stmt="quantize_to_fp4(x, config)",
globals={"x": x, "config": config},
)
print(f"{t.timeit(repeats).mean * 1000:.4f}ms")
@app.function(image=img, cpu=4, memory=8 * 1024, gpu="B200")
def run_speedtest_on_modal(**kwargs: dict[str, Any]) -> None:
run_speedtest(**kwargs)
@click.command()
@click.option("--block-scale-2d", is_flag=True)
@click.option("--input-shape", type=str, default="1024,1024")
@click.option("--modal", is_flag=True)
@click.option("--repeats", type=int, default=100)
@click.option("--rht", is_flag=True)
@click.option("--round-style", type=RoundStyle, default=RoundStyle.nearest)
@click.option("--scale-rule", type=ScaleRule, default=ScaleRule.mse)
@click.option("--transpose", is_flag=True)
def cli(**kwargs: dict[str, Any]) -> None:
if kwargs.pop("modal"):
with modal.enable_output(), app.run():
run_speedtest_on_modal.remote(**kwargs)
else:
run_speedtest(**kwargs)
if __name__ == "__main__":
cli()
fouroversix/scripts/test_on_modal.py
+30
View File
@@ -0,0 +1,30 @@
from __future__ import annotations
from pathlib import Path
from .resources import FOUROVERSIX_INSTALL_PATH, Dependency, app, get_image
img = get_image(
dependencies=[Dependency.transformer_engine, Dependency.fouroversix],
include_tests=True,
)
with img.imports():
import pytest
@app.function(image=img, cpu=4, memory=8 * 1024, gpu="B200", timeout=30 * 60)
def run_tests(*args: list[str]) -> None:
"""Run tests on a B200 on Modal."""
args = list(args)
tests_path = (Path(FOUROVERSIX_INSTALL_PATH) / "tests").as_posix()
if len(args) == 0:
args = [tests_path]
elif "tests" in args[0] or "test_" in args[0]:
args[0] = (Path(FOUROVERSIX_INSTALL_PATH) / args[0]).as_posix()
else:
args = [tests_path, *args]
pytest.main(args)
fouroversix/scripts/train/__init__.py
View File
fouroversix/scripts/train/__main__.py
+244
View File
@@ -0,0 +1,244 @@
from __future__ import annotations
import os
import subprocess
import time
from pathlib import Path
from typing import Any
import click
import modal
from ..resources import (
FOUROVERSIX_CACHE_PATH,
FOUROVERSIX_INSTALL_PATH,
Dependency,
app,
cache_volume,
get_image,
wandb_secret,
)
img = get_image(
dependencies=[Dependency.flash_attention, Dependency.flame, Dependency.fouroversix],
extra_pip_dependencies=["datasets<4.6"],
)
def train(
*,
batch_size: int,
checkpoint_interval: int,
checkpoint_keep_latest_k: int,
checkpoint_load_step: int,
context_length: int,
dataset: str,
dataset_name: str,
dataset_split: str,
exp_folder: str,
gradient_accumulation_steps: int,
initial_load_path: str | None,
job_config_file: str,
lr: float,
lr_decay_type: str,
model_config: str,
model_name: str,
no_torch_compile: bool,
seed: int,
tokenizer: str,
training_steps: int | None,
) -> None:
import torch
# Cache activations and gradients and set dump folder
os.environ["CACHE_ACTIVATIONS"] = "1"
os.environ["CACHE_GRADIENTS"] = "1"
os.environ["DUMP_FOLDER"] = f"{exp_folder}/{model_name}"
# Set MODEL_NAME for wandb
os.environ["MODEL_NAME"] = model_name
# Set NGPU for flame
os.environ["NGPU"] = str(torch.cuda.device_count())
if training_steps is None:
if dataset_name == "sample-10BT":
num_tokens = 10_000_000_000
elif dataset_name == "sample-100BT":
num_tokens = 100_000_000_000
else:
msg = (
"You must provide the number of training steps if not using the "
"sample-10BT or sample-100BT datasets"
)
raise ValueError(msg)
training_steps = num_tokens // int(
context_length * batch_size * torch.cuda.device_count(),
)
# Start training
args = [
"bash",
"train.sh",
"--job.config_file",
job_config_file,
"--job.dump_folder",
f"{exp_folder}/{model_name}",
"--model.config",
model_config,
"--model.tokenizer_path",
tokenizer,
"--optimizer.name",
"AdamW",
"--optimizer.lr",
str(lr),
"--lr_scheduler.warmup_steps",
"0",
"--lr_scheduler.decay_ratio",
"0.15",
"--lr_scheduler.decay_type",
lr_decay_type,
"--lr_scheduler.lr_min",
"0.01",
"--training.batch_size",
"1",
"--training.seq_len",
str(int(context_length * batch_size)),
"--training.context_len",
str(context_length),
"--training.varlen",
"--training.gradient_accumulation_steps",
str(gradient_accumulation_steps),
"--training.steps",
str(training_steps),
"--training.max_norm",
"1.0",
"--training.skip_nan_inf",
"--training.dataset",
dataset,
"--training.dataset_name",
dataset_name,
"--training.dataset_split",
dataset_split,
"--training.num_workers",
"32",
"--training.prefetch_factor",
"2",
"--training.seed",
str(seed),
"--checkpoint.interval",
str(checkpoint_interval),
"--checkpoint.load_step",
str(checkpoint_load_step),
"--checkpoint.keep_latest_k",
str(checkpoint_keep_latest_k),
"--metrics.log_freq",
"1",
]
if not no_torch_compile:
args.append("--training.compile")
if initial_load_path is not None:
args.extend(
[
"--checkpoint.initial_load_path",
initial_load_path,
"--checkpoint.no_initial_load_model_weights_only",
],
)
subprocess.run(args, check=True)
@app.cls(
image=img,
gpu="B200:8",
timeout=24 * 60 * 60,
cpu=64,
memory=8 * 64 * 1024,
volumes={FOUROVERSIX_CACHE_PATH: cache_volume},
secrets=[wandb_secret],
)
class ModalTrainer:
"""Run training jobs on Modal."""
@modal.method()
def train(self, **kwargs: dict[str, Any]) -> None:
"""Start a training job on Modal."""
os.chdir(FOUROVERSIX_INSTALL_PATH / "third_party" / "flame")
train(**kwargs)
@click.command()
@click.option("--batch-size", type=float, default=16)
@click.option("--checkpoint-interval", type=int, default=1000)
@click.option("--checkpoint-keep-latest-k", type=int, default=0)
@click.option("--checkpoint-load-step", type=int, default=-1)
@click.option("--context-length", type=int, default=8192)
@click.option("--dataset", type=str, default="HuggingFaceFW/fineweb-edu")
@click.option("--dataset-name", type=str, default="sample-100BT")
@click.option("--dataset-split", type=str, default="train")
@click.option("--detach", is_flag=True)
@click.option("--exp-folder", type=str, default="exp")
@click.option("--gradient-accumulation-steps", type=int, default=1)
@click.option("--initial-load-path", type=str)
@click.option("--job-config-file", type=str, default="flame/models/fla.toml")
@click.option("--lr", type=float, default=1.2e-3)
@click.option("--lr-decay-type", type=str, default="linear")
@click.option("--modal", is_flag=True)
@click.option("--modal-gpu", type=str, default="B200:8")
@click.option("--model-config", type=str, required=True)
@click.option("--model-name", type=str, required=True)
@click.option("--no-torch-compile", is_flag=True)
@click.option("--seed", type=int, default=42)
@click.option("--tokenizer", type=str, default="fla-hub/transformer-1.3B-100B")
@click.option("--training-steps", type=int, default=None)
@click.option("--wait-for-pid", type=int, default=None)
def cli(**kwargs: dict[str, Any]) -> None:
# Options that are not passed to the train function
detach = kwargs.pop("detach", False)
modal_gpu = kwargs.pop("modal_gpu", "B200:8")
use_modal = kwargs.pop("modal", False)
wait_for_pid = kwargs.pop("wait_for_pid", None)
# Wait for the previous training job to finish
if wait_for_pid is not None:
while (
subprocess.run(["kill", "-0", str(wait_for_pid)], check=False).returncode
== 0
):
time.sleep(1)
time.sleep(60)
if not Path(kwargs["model_config"]).exists():
kwargs["model_config"] = (
Path(__file__).parent.parent.parent / kwargs["model_config"]
)
if not Path(kwargs["model_config"]).exists():
msg = f"Model config file not found: {kwargs['model_config']}"
raise FileNotFoundError(msg)
# Set exp folder on Modal
if use_modal:
with modal.enable_output(), app.run(detach=detach):
kwargs["exp_folder"] = (FOUROVERSIX_CACHE_PATH / "exp").as_posix()
kwargs["model_config"] = (
(FOUROVERSIX_INSTALL_PATH / kwargs["model_config"])
.absolute()
.as_posix()
)
ModalTrainer.with_options(gpu=modal_gpu)().train.remote(**kwargs)
else:
kwargs["model_config"] = Path(kwargs["model_config"]).absolute().as_posix()
os.chdir(Path(__file__).parent.parent.parent / "third_party" / "flame")
train(**kwargs)
if __name__ == "__main__":
cli()
fouroversix/scripts/train/configs/transformer_1B_bf16.json
+26
View File
@@ -0,0 +1,26 @@
{
"attention_bias": false,
"bos_token_id": 1,
"eos_token_id": 2,
"fuse_cross_entropy": true,
"fuse_norm": true,
"fuse_swiglu": false,
"hidden_act": "swish",
"hidden_ratio": 4,
"hidden_size": 2048,
"initializer_range": 0.02,
"max_position_embeddings": 8192,
"model_type": "fp4_transformer",
"num_heads": 32,
"num_hidden_layers": 24,
"norm_eps": 1e-06,
"tie_word_embeddings": false,
"use_cache": true,
"vocab_size": 32000,
"qk_norm": true,
"layer_precision_configs": [
{
"repeats": 24
}
]
}
fouroversix/scripts/train/configs/transformer_1B_fp4.json
+31
View File
@@ -0,0 +1,31 @@
{
"attention_bias": false,
"bos_token_id": 1,
"eos_token_id": 2,
"fuse_cross_entropy": true,
"fuse_norm": true,
"fuse_swiglu": false,
"hidden_act": "swish",
"hidden_ratio": 4,
"hidden_size": 2048,
"initializer_range": 0.02,
"max_position_embeddings": 8192,
"model_type": "fp4_transformer",
"num_heads": 32,
"num_hidden_layers": 24,
"norm_eps": 1e-06,
"tie_word_embeddings": false,
"use_cache": true,
"vocab_size": 32000,
"qk_norm": true,
"layer_precision_configs": [
{
"repeats": 20,
"dtype": "nvfp4",
"scale_rule": "static_6"
},
{
"repeats": 4
}
]
}
fouroversix/scripts/train/configs/transformer_1B_fp4_fouroversix.json
+31
View File
@@ -0,0 +1,31 @@
{
"attention_bias": false,
"bos_token_id": 1,
"eos_token_id": 2,
"fuse_cross_entropy": true,
"fuse_norm": true,
"fuse_swiglu": false,
"hidden_act": "swish",
"hidden_ratio": 4,
"hidden_size": 2048,
"initializer_range": 0.02,
"max_position_embeddings": 8192,
"model_type": "fp4_transformer",
"num_heads": 32,
"num_hidden_layers": 24,
"norm_eps": 1e-06,
"tie_word_embeddings": false,
"use_cache": true,
"vocab_size": 32000,
"qk_norm": true,
"layer_precision_configs": [
{
"repeats": 20,
"dtype": "nvfp4",
"scale_rule": "mse"
},
{
"repeats": 4
}
]
}
fouroversix/scripts/train/configs/transformer_340M_bf16.json
+25
View File
@@ -0,0 +1,25 @@
{
"attention_bias": false,
"bos_token_id": 1,
"eos_token_id": 2,
"fuse_cross_entropy": true,
"fuse_norm": true,
"fuse_swiglu": false,
"hidden_act": "swish",
"hidden_size": 1024,
"initializer_range": 0.02,
"max_position_embeddings": 8192,
"model_type": "fp4_transformer",
"num_heads": 16,
"num_hidden_layers": 24,
"norm_eps": 1e-06,
"tie_word_embeddings": false,
"use_cache": true,
"vocab_size": 32000,
"qk_norm": true,
"layer_precision_configs": [
{
"repeats": 24
}
]
}
fouroversix/scripts/train/configs/transformer_340M_fp4.json
+30
View File
@@ -0,0 +1,30 @@
{
"attention_bias": false,
"bos_token_id": 1,
"eos_token_id": 2,
"fuse_cross_entropy": true,
"fuse_norm": true,
"fuse_swiglu": false,
"hidden_act": "swish",
"hidden_size": 1024,
"initializer_range": 0.02,
"max_position_embeddings": 8192,
"model_type": "fp4_transformer",
"num_heads": 16,
"num_hidden_layers": 24,
"norm_eps": 1e-06,
"tie_word_embeddings": false,
"use_cache": true,
"vocab_size": 32000,
"qk_norm": true,
"layer_precision_configs": [
{
"repeats": 20,
"dtype": "nvfp4",
"scale_rule": "static_6"
},
{
"repeats": 4
}
]
}
fouroversix/scripts/train/configs/transformer_340M_fp4_fouroversix.json
+30
View File
@@ -0,0 +1,30 @@
{
"attention_bias": false,
"bos_token_id": 1,
"eos_token_id": 2,
"fuse_cross_entropy": true,
"fuse_norm": true,
"fuse_swiglu": false,
"hidden_act": "swish",
"hidden_size": 1024,
"initializer_range": 0.02,
"max_position_embeddings": 8192,
"model_type": "fp4_transformer",
"num_heads": 16,
"num_hidden_layers": 24,
"norm_eps": 1e-06,
"tie_word_embeddings": false,
"use_cache": true,
"vocab_size": 32000,
"qk_norm": true,
"layer_precision_configs": [
{
"repeats": 20,
"dtype": "nvfp4",
"scale_rule": "mse"
},
{
"repeats": 4
}
]
}
fouroversix/scripts/train/prepare_dataset.py
+15
View File
@@ -0,0 +1,15 @@
from ..resources import FOUROVERSIX_CACHE_PATH, app, cache_volume, get_image
img = get_image(dependencies=[], extra_pip_dependencies=["datasets"])
with img.imports():
from datasets import load_dataset
@app.function(
image=img,
timeout=24 * 60 * 60,
volumes={FOUROVERSIX_CACHE_PATH: cache_volume},
)
def prepare_dataset(path: str, name: str) -> None:
load_dataset(path, name)
fouroversix/setup.py
+296
View File
@@ -0,0 +1,296 @@
import functools
import os
import platform
import subprocess
import sys
import urllib.request
import warnings
from pathlib import Path
from typing import Any
import torch
from packaging.version import Version, parse
from setuptools import setup
from setuptools.command.bdist_wheel import bdist_wheel
from torch.utils.cpp_extension import CUDA_HOME, BuildExtension, CUDAExtension
BASE_WHEEL_URL = "https://github.com/mit-han-lab/fouroversix/releases/download"
PACKAGE_NAME = "fouroversix"
PACKAGE_VERSION = "1.1.0"
CUTLASS_DEBUG = os.getenv("CUTLASS_DEBUG", "0") == "1"
FORCE_BUILD = os.getenv("FORCE_BUILD", "0") == "1"
FORCE_CXX11_ABI = os.getenv("FORCE_CXX11_ABI", "0") == "1"
SKIP_CUDA_BUILD = os.getenv("SKIP_CUDA_BUILD", "0") == "1"
@functools.cache
def get_cuda_archs() -> list[str]:
return os.getenv("CUDA_ARCHS", "100;103;110;120").split(";")
def get_cuda_bare_metal_version() -> Version | None:
if CUDA_HOME is None:
warnings.warn(
"nvcc was not found. Are you sure your environment has nvcc available? If "
"you're installing within a container from "
"https://hub.docker.com/r/pytorch/pytorch, only images with 'devel' in "
"their name will provide nvcc.",
stacklevel=1,
)
return None
raw_output = subprocess.check_output(
[CUDA_HOME + "/bin/nvcc", "-V"],
universal_newlines=True,
)
output = raw_output.split()
release_idx = output.index("release") + 1
return parse(output[release_idx].split(",")[0])
def get_cuda_gencodes() -> list[str]:
"""
Add -gencode flags based on nvcc capabilities.
Uses the following rules:
- sm_100/120 on CUDA >= 12.8
- Use 100f on CUDA >= 12.9 (Blackwell family-specific)
- Map requested 110 -> 101 if CUDA < 13.0 (Thor rename)
- Embed PTX for newest arch for forward compatibility
"""
archs = set(get_cuda_archs())
cuda_version = get_cuda_bare_metal_version()
cc_flags = []
# Blackwell requires >= 12.8
if cuda_version is not None and cuda_version >= Version("12.8"):
if "100" in archs:
cc_flags += ["-gencode", "arch=compute_100a,code=sm_100a"]
if "103" in archs:
cc_flags += ["-gencode", "arch=compute_103a,code=sm_103a"]
# Thor rename: 12.9 uses sm_101; 13.0+ uses sm_110
if "110" in archs:
if cuda_version >= Version("13.0"):
cc_flags += ["-gencode", "arch=compute_110f,code=sm_110"]
elif cuda_version >= Version("12.9"):
# Provide Thor support for CUDA 12.9 via sm_101
cc_flags += ["-gencode", "arch=compute_101f,code=sm_101"]
# else: no Thor support in older toolkits
if "120" in archs:
# sm_120 is supported in CUDA 12.8/12.9+ toolkits
if cuda_version >= Version("12.9"):
cc_flags += ["-gencode", "arch=compute_120f,code=sm_120"]
else:
cc_flags += ["-gencode", "arch=compute_120a,code=sm_120a"]
return cc_flags
def get_platform() -> str:
if sys.platform.startswith("linux"):
return f"linux_{platform.uname().machine}"
if sys.platform == "darwin":
mac_version = ".".join(platform.mac_ver()[0].split(".")[:2])
return f"macosx_{mac_version}_x86_64"
if sys.platform == "win32":
return "win_amd64"
msg = f"Unsupported platform: {sys.platform}"
raise ValueError(msg)
def get_wheel_url() -> tuple[str, str]:
torch_version_raw = parse(torch.__version__)
python_version = f"cp{sys.version_info.major}{sys.version_info.minor}"
platform_name = get_platform()
torch_version = f"{torch_version_raw.major}.{torch_version_raw.minor}"
cxx11_abi = str(torch._C._GLIBCXX_USE_CXX11_ABI).upper() # noqa: SLF001
# We only compile for CUDA 12.8 to save CI time. Minor versions should be
# compatible.
torch_cuda_version = parse("12.8")
cuda_version = f"cu{torch_cuda_version.major}"
wheel_filename = (
f"{PACKAGE_NAME}-{PACKAGE_VERSION}+{cuda_version}torch{torch_version}"
f"cxx11abi{cxx11_abi}-{python_version}-{python_version}-{platform_name}.whl"
)
return f"{BASE_WHEEL_URL}/v{PACKAGE_VERSION}/{wheel_filename}", wheel_filename
class CachedWheelsCommand(bdist_wheel):
"""
Custom bdist wheel command that checks for pre-built wheels on GitHub Releases.
The CachedWheelsCommand plugs into the default bdist wheel, which is ran by pip
when it cannot find an existing wheel (which is currently the case for all
fouroversix installs). We use the environment parameters to detect whether there is
already a pre-built version of a compatible wheel available and short-circuits the
standard full build pipeline.
Credit: https://github.com/Dao-AILab/flash-attention/blob/main/setup.py
"""
def run(self) -> None:
"""Run the command."""
if FORCE_BUILD:
return super().run()
wheel_url, wheel_filename = get_wheel_url()
print(f"Guessing wheel URL: {wheel_url}")
try:
urllib.request.urlretrieve(wheel_url, wheel_filename) # noqa: S310
# Make the archive
# Lifted from the root wheel processing command
# https://github.com/pypa/wheel/blob/cf71108ff9f6ffc36978069acb28824b44ae028e/src/wheel/bdist_wheel.py#LL381C9-L381C85
if not Path(self.dist_dir).exists():
Path(self.dist_dir).mkdir(parents=True, exist_ok=True)
impl_tag, abi_tag, plat_tag = self.get_tag()
archive_basename = f"{self.wheel_dist_name}-{impl_tag}-{abi_tag}-{plat_tag}"
wheel_path = Path(self.dist_dir) / (archive_basename + ".whl")
print(f"Raw wheel path: {wheel_path}")
Path(wheel_filename).rename(wheel_path)
except (urllib.error.HTTPError, urllib.error.URLError):
print("Precompiled wheel not found. Building from source...")
# If the wheel could not be downloaded, build from source
super().run()
class NinjaBuildExtension(BuildExtension):
"""
Custom build extension that tells Ninja how many jobs to run.
Credit: https://github.com/Dao-AILab/flash-attention/blob/main/setup.py
"""
def __init__(self, *args: list[Any], **kwargs: dict[str, Any]) -> None:
# do not override env MAX_JOBS if already exists
if not os.environ.get("MAX_JOBS"):
try:
import psutil
# calculate the maximum allowed NUM_JOBS based on cores
max_num_jobs_cores = max(1, os.cpu_count() // 2)
# calculate the maximum allowed NUM_JOBS based on free memory
free_memory_gb = psutil.virtual_memory().available / (
1024**3
) # free memory in GB
max_num_jobs_memory = int(
free_memory_gb / 9,
) # each JOB peak memory cost is ~8-9GB when threads = 4
# pick lower value of jobs based on cores vs memory metric to minimize
# oom and swap usage during compilation
max_jobs = max(1, min(max_num_jobs_cores, max_num_jobs_memory))
os.environ["MAX_JOBS"] = str(max_jobs)
except ImportError:
warnings.warn(
"psutil not found, install psutil and ninja to get better build "
"performance",
stacklevel=1,
)
super().__init__(*args, **kwargs)
if SKIP_CUDA_BUILD:
warnings.warn(
"SKIP_CUDA_BUILD is set to 1, installing fouroversix without quantization and "
"matmul kernels",
stacklevel=1,
)
ext_modules = None
else:
if Path(".git").exists():
subprocess.run(
[
"git",
"submodule",
"update",
"--init",
"third_party/cutlass",
],
check=True,
)
elif not Path("third_party/cutlass").exists():
msg = (
"third_party/cutlass is missing, please use source distribution or git "
"clone"
)
raise RuntimeError(msg)
# The compiler flag -D_GLIBCXX_USE_CXX11_ABI is set to be the same as
# torch._C._GLIBCXX_USE_CXX11_ABI
# https://github.com/pytorch/pytorch/blob/8472c24e3b5b60150096486616d98b7bea01500b/torch/utils/cpp_extension.py#L920
if FORCE_CXX11_ABI:
torch._C._GLIBCXX_USE_CXX11_ABI = True # noqa: SLF001
setup_dir = Path(__file__).parent
kernels_dir = setup_dir / "src" / "fouroversix" / "csrc"
sources = [
path.relative_to(Path(__file__).parent).as_posix()
for ext in ["**/*.cu", "**/*.cpp"]
for path in kernels_dir.glob(ext)
]
cxx_compile_args = ["-std=c++17"]
nvcc_compile_args = [
"-std=c++17",
"--expt-relaxed-constexpr",
"-Xcompiler",
"-funroll-loops",
"-Xcompiler",
"-finline-functions",
*get_cuda_gencodes(),
]
if CUTLASS_DEBUG:
nvcc_compile_args.extend(
[
"-O0",
"-DCUTLASS_DEBUG_TRACE_LEVEL=3",
"-DCUTLASS_DEBUG_ENABLE=1",
"-g",
],
)
else:
cxx_compile_args.extend(["-O3"])
nvcc_compile_args.extend(["-O3", "-DNDEBUG"])
ext_modules = [
CUDAExtension(
"fouroversix._C",
sources,
extra_compile_args={"cxx": cxx_compile_args, "nvcc": nvcc_compile_args},
include_dirs=[
setup_dir / "third_party/cutlass/examples/common",
setup_dir / "third_party/cutlass/include",
setup_dir / "third_party/cutlass/tools/util/include",
kernels_dir / "include",
],
),
]
setup(
name=PACKAGE_NAME,
version=PACKAGE_VERSION,
ext_modules=ext_modules,
cmdclass={
"bdist_wheel": CachedWheelsCommand,
"build_ext": NinjaBuildExtension,
},
include_package_data=True,
)
fouroversix/src/fouroversix/__init__.py
+33
View File
@@ -0,0 +1,33 @@
from importlib.metadata import version
from .matmul import fp4_matmul
from .model import (
FourOverSixLinear,
ModelQuantizationConfig,
ModuleQuantizationConfig,
QuantizedModule,
quantize_model,
)
from .quantize import QuantizationConfig, QuantizedTensor, quantize_to_fp4
from .utils import DataType, MatmulBackend, QuantizeBackend, RoundStyle, ScaleRule
from .weight_conversions import WeightConversions
__version__ = version("fouroversix")
__all__ = [
"DataType",
"FourOverSixLinear",
"MatmulBackend",
"ModelQuantizationConfig",
"ModuleQuantizationConfig",
"QuantizationConfig",
"QuantizeBackend",
"QuantizedModule",
"QuantizedTensor",
"RoundStyle",
"ScaleRule",
"WeightConversions",
"fp4_matmul",
"quantize_model",
"quantize_to_fp4",
]
fouroversix/src/fouroversix/csrc/bindings.cpp
+36
View File
@@ -0,0 +1,36 @@
#include <Python.h>
#include <torch/extension.h>
extern "C"
{
PyObject *PyInit__C(void)
{
static struct PyModuleDef module_def = {
PyModuleDef_HEAD_INIT,
"_C", /* name of module */
NULL, /* module documentation, may be NULL */
-1, /* size of per-interpreter state of the module,
or -1 if the module keeps state in global variables. */
NULL, /* methods */
NULL,
NULL,
NULL,
NULL,
};
return PyModule_Create(&module_def);
}
}
namespace fouroversix
{
TORCH_LIBRARY(fouroversix, m)
{
m.def("quantize_to_fp4(Tensor x, bool is_nvfp4, bool is_rtn, bool is_rht, bool is_2d, bool is_transpose, int selection_rule, int rbits) -> (Tensor, Tensor, Tensor)");
m.def("gemm_mxfp4mxfp4_accum_fp32_out_bf16_tnt(Tensor A, Tensor B, Tensor A_sf, Tensor B_sf, Tensor alpha) -> Tensor");
m.def("gemm_mxfp4mxfp4_accum_fp32_out_bf16_tnt_sm120(Tensor A, Tensor B, Tensor A_sf, Tensor B_sf, Tensor alpha) -> Tensor");
m.def("gemm_nvfp4nvfp4_accum_fp32_out_bf16_tnt(Tensor A, Tensor B, Tensor A_sf, Tensor B_sf, Tensor alpha) -> Tensor");
m.def("gemm_nvfp4nvfp4_accum_fp32_out_bf16_tnt_sm120(Tensor A, Tensor B, Tensor A_sf, Tensor B_sf, Tensor alpha) -> Tensor");
m.def("gemm_nvfp4nvfp4_accum_fp32_out_fp16_tnt(Tensor A, Tensor B, Tensor A_sf, Tensor B_sf, Tensor alpha) -> Tensor");
m.def("gemm_nvfp4nvfp4_accum_fp32_out_fp16_tnt_sm120(Tensor A, Tensor B, Tensor A_sf, Tensor B_sf, Tensor alpha) -> Tensor");
}
}
fouroversix/src/fouroversix/csrc/fp4_gemm.cu
+187
View File
@@ -0,0 +1,187 @@
#include "cutlass/cutlass.h"
#include "cute/tensor.hpp"
#include "cutlass/tensor_ref.h"
#include "cutlass/epilogue/thread/linear_combination.h"
#include "cutlass/gemm/dispatch_policy.hpp"
#include "cutlass/gemm/collective/collective_builder.hpp"
#include "cutlass/epilogue/collective/collective_builder.hpp"
#include "cutlass/detail/sm100_blockscaled_layout.hpp"
#include "cutlass/gemm/device/gemm_universal_adapter.h"
#include "cutlass/gemm/kernel/gemm_universal.hpp"
#include "cutlass/gemm/kernel/tile_scheduler_params.h"
#include "cutlass/util/command_line.h"
#include "cutlass/util/distribution.h"
#include "cutlass/util/host_tensor.h"
#include "cutlass/util/packed_stride.hpp"
#include "cutlass/util/tensor_view_io.h"
#include "cutlass/util/reference/device/gemm.h"
#include "cutlass/util/reference/device/tensor_compare.h"
#include "cutlass/util/reference/host/tensor_fill.h"
#include "cutlass/util/reference/host/gett.hpp"
#include "cutlass/util/reference/host/tensor_norm.h"
#include "cutlass/util/reference/host/tensor_compare.h"
#include "helper.h"
#include <ATen/ATen.h>
#include <ATen/cuda/CUDAContext.h>
#include <torch/all.h>
#include <c10/cuda/CUDAGuard.h>
#include "element_traits.hpp"
namespace fouroversix
{
using namespace cute;
// Adapted from example 72b
template <typename ElementA,
typename MmaTileShape = Shape<_128, _128, _256>,
typename ClusterShape = Shape<_2, _4, _1>,
typename KernelMainloopPolicy = cutlass::gemm::collective::KernelScheduleAuto,
int AlignmentA = 32,
typename ElementB = ElementA,
int AlignmentB = 32,
typename LayoutATag = cutlass::layout::RowMajor,
typename LayoutBTag = cutlass::layout::ColumnMajor,
typename ElementD = cutlass::bfloat16_t,
typename ArchTag = cutlass::arch::Sm100>
torch::Tensor gemm_fp4fp4_accum_fp32(torch::Tensor const &A, torch::Tensor const &B, torch::Tensor const &A_sf, torch::Tensor const &B_sf, torch::Tensor const &alpha)
{
at::cuda::CUDAGuard device_guard(A.device());
// C/D matrix configuration
using ElementC = void; // Element type for C matrix operand
using LayoutCTag = cutlass::layout::RowMajor; // Layout type for C matrix operand
using LayoutDTag = cutlass::layout::RowMajor; // Layout type for D matrix operand
constexpr int AlignmentC = 1; // Memory access granularity/alignment of C matrix in units of elements (up to 16 bytes)
constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value; // Memory access granularity/alignment of D matrix in units of elements (up to 16 bytes)
// Kernel functional config
using ElementAccumulator = float; // Element type for internal accumulation
using OperatorClass = cutlass::arch::OpClassBlockScaledTensorOp; // Operator class tag
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
ArchTag, OperatorClass,
MmaTileShape, ClusterShape,
cutlass::epilogue::collective::EpilogueTileAuto,
ElementAccumulator, ElementAccumulator,
ElementC, LayoutCTag, AlignmentC,
ElementD, LayoutDTag, AlignmentD,
cutlass::epilogue::collective::EpilogueScheduleAuto // Epilogue schedule policy
>::CollectiveOp;
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
ArchTag, OperatorClass,
ElementA, LayoutATag, AlignmentA,
ElementB, LayoutBTag, AlignmentB,
ElementAccumulator,
MmaTileShape, ClusterShape,
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
KernelMainloopPolicy>::CollectiveOp;
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
Shape<int, int, int, int>, // Indicates ProblemShape
CollectiveMainloop,
CollectiveEpilogue,
void>;
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
// Reference device GEMM implementation type
using StrideA = typename Gemm::GemmKernel::StrideA;
using StrideB = typename Gemm::GemmKernel::StrideB;
using StrideC = typename Gemm::GemmKernel::StrideC;
using StrideD = typename Gemm::GemmKernel::StrideD;
using LayoutSFA = typename Gemm::GemmKernel::CollectiveMainloop::LayoutSFA; // Scale Factor tensors have an interleaved layout. Bring Layout instead of stride.
using LayoutSFB = typename Gemm::GemmKernel::CollectiveMainloop::LayoutSFB; // Scale Factor tensors have an interleaved layout. Bring Layout instead of stride.
using Sm1xxBlkScaledConfig = typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig;
torch::checkAllContiguous("gemm_fp4fp4_accum_fp32_out_bf16", {{A, "A", 0}, {B, "B", 1}, {A_sf, "A_sf", 2}, {B_sf, "B_sf", 3}, {alpha, "alpha", 4}});
torch::checkDeviceType("gemm_fp4fp4_accum_fp32_out_bf16", {A, B, A_sf, B_sf, alpha}, at::DeviceType::CUDA);
torch::checkAllSameGPU("gemm_fp4fp4_accum_fp32_out_bf16", {{A, "A", 0}, {B, "B", 1}, {A_sf, "A_sf", 2}, {B_sf, "B_sf", 3}, {alpha, "alpha", 4}});
check_block_scale_factor_type<ElementA>(A_sf, "A_sf");
check_block_scale_factor_type<ElementB>(B_sf, "B_sf");
auto [M, N, K] = check_and_get_fp4_matmul_dims<ElementA, LayoutATag, ElementB, LayoutBTag>(A, B, A_sf, B_sf);
auto D = torch::empty({M, N}, torch::dtype(element_traits<ElementD>::scalar_type).device(A.device()));
Gemm gemm;
// Create stride and layout information for the packed tensors
// For packed NVFP4 tensors, we need to use the appropriate stride and layout
StrideA stride_A = cutlass::make_cute_packed_stride(StrideA{}, {M, K, 1});
StrideB stride_B = cutlass::make_cute_packed_stride(StrideB{}, {N, K, 1});
StrideC stride_C = cutlass::make_cute_packed_stride(StrideC{}, {M, N, 1});
StrideD stride_D = cutlass::make_cute_packed_stride(StrideD{}, {M, N, 1});
// Create scale factor layouts
LayoutSFA layout_SFA = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(make_shape(M, N, K, 1));
LayoutSFB layout_SFB = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(make_shape(M, N, K, 1));
typename Gemm::Arguments args{
cutlass::gemm::GemmUniversalMode::kGemm,
{M, N, K, 1},
{// Mainloop arguments
static_cast<typename ElementA::DataType const *>(A.data_ptr()), stride_A,
static_cast<typename ElementB::DataType const *>(B.data_ptr()), stride_B,
static_cast<typename ElementA::ScaleFactorType const *>(A_sf.data_ptr()), layout_SFA,
static_cast<typename ElementB::ScaleFactorType const *>(B_sf.data_ptr()), layout_SFB},
{// Epilogue arguments
{1.0f, 0.0f},
nullptr,
stride_C,
static_cast<ElementD *>(D.data_ptr()),
stride_D}};
args.epilogue.thread.alpha_ptr = static_cast<ElementAccumulator *>(alpha.data_ptr());
// Check if the problem size is supported or not
CUTLASS_CHECK(gemm.can_implement(args));
// Initialize CUTLASS kernel with arguments and workspace pointer
CUTLASS_CHECK(gemm.initialize(args));
CUTLASS_CHECK(gemm.run(args));
return D;
}
TORCH_LIBRARY_IMPL(fouroversix, CUDA, m)
{
/* MXFP4 */
m.impl("gemm_mxfp4mxfp4_accum_fp32_out_bf16_tnt",
&gemm_fp4fp4_accum_fp32<
cutlass::mx_float4_t<cutlass::float_e2m1_t>,
Shape<_256, _256, _256>,
Shape<_4, _1, _1>,
cutlass::gemm::KernelTmaWarpSpecialized2SmMxf4Sm100>);
/* NVFP4 */
m.impl("gemm_nvfp4nvfp4_accum_fp32_out_bf16_tnt",
&gemm_fp4fp4_accum_fp32<
cutlass::nv_float4_t<cutlass::float_e2m1_t>,
Shape<_256, _256, _256>,
Shape<_4, _1, _1>,
cutlass::gemm::KernelTmaWarpSpecialized2SmNvf4Sm100>);
m.impl("gemm_nvfp4nvfp4_accum_fp32_out_fp16_tnt",
&gemm_fp4fp4_accum_fp32<
cutlass::nv_float4_t<cutlass::float_e2m1_t>,
Shape<_256, _256, _256>,
Shape<_4, _1, _1>,
cutlass::gemm::KernelTmaWarpSpecialized2SmNvf4Sm100,
32,
cutlass::nv_float4_t<cutlass::float_e2m1_t>,
32,
cutlass::layout::RowMajor,
cutlass::layout::ColumnMajor,
cutlass::half_t>);
}
}
fouroversix/src/fouroversix/csrc/fp4_gemm_sm120.cu
+414
View File
@@ -0,0 +1,414 @@
#include "cutlass/cutlass.h"
#include "cute/tensor.hpp"
#include "cutlass/tensor_ref.h"
#include "cutlass/epilogue/thread/linear_combination.h"
#include "cutlass/gemm/dispatch_policy.hpp"
#include "cutlass/gemm/collective/collective_builder.hpp"
#include "cutlass/epilogue/collective/collective_builder.hpp"
#include "cutlass/detail/sm100_blockscaled_layout.hpp"
#include "cutlass/gemm/device/gemm_universal_adapter.h"
#include "cutlass/gemm/kernel/gemm_universal.hpp"
#include "cutlass/gemm/kernel/tile_scheduler_params.h"
#include "cutlass/util/command_line.h"
#include "cutlass/util/distribution.h"
#include "cutlass/util/host_tensor.h"
#include "cutlass/util/packed_stride.hpp"
#include "cutlass/util/tensor_view_io.h"
#include "cutlass/util/reference/device/gemm.h"
#include "cutlass/util/reference/device/tensor_compare.h"
#include "cutlass/util/reference/host/tensor_fill.h"
#include "cutlass/util/reference/host/gett.hpp"
#include "cutlass/util/reference/host/tensor_norm.h"
#include "cutlass/util/reference/host/tensor_compare.h"
#include "helper.h"
#include <ATen/ATen.h>
#include <ATen/cuda/CUDAContext.h>
#include <torch/all.h>
#include <c10/cuda/CUDAGuard.h>
#include "element_traits.hpp"
// NOTE: There seems to be an issue with NVCC that causes runtime errors in SM120 GEMMs
// when they are compiled as templated functions. As a result, this file contains only
// non-templated functions. See: https://github.com/NVIDIA/cutlass/issues/2478
namespace fouroversix
{
using namespace cute;
// Adapted from example 72b
torch::Tensor gemm_mxfp4mxfp4_accum_fp32_out_bf16_tnt_sm120(torch::Tensor const &A, torch::Tensor const &B, torch::Tensor const &A_sf, torch::Tensor const &B_sf, torch::Tensor const &alpha)
{
at::cuda::CUDAGuard device_guard(A.device());
using ElementA = cutlass::mx_float4_t<cutlass::float_e2m1_t>;
using ElementB = cutlass::mx_float4_t<cutlass::float_e2m1_t>;
using ElementD = cutlass::bfloat16_t;
using MmaTileShape = Shape<_128, _128, _128>;
using ClusterShape = Shape<_1, _1, _1>;
// C/D matrix configuration
using ElementC = void; // Element type for C matrix operand
using LayoutATag = cutlass::layout::RowMajor; // Layout type for A matrix operand
using LayoutBTag = cutlass::layout::ColumnMajor; // Layout type for B matrix operand
using LayoutCTag = cutlass::layout::RowMajor; // Layout type for C matrix operand
using LayoutDTag = cutlass::layout::RowMajor; // Layout type for D matrix operand
constexpr int AlignmentA = 32;
constexpr int AlignmentB = 32;
constexpr int AlignmentC = 1; // Memory access granularity/alignment of C matrix in units of elements (up to 16 bytes)
constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value; // Memory access granularity/alignment of D matrix in units of elements (up to 16 bytes)
// Kernel functional config
using ElementAccumulator = float; // Element type for internal accumulation
using ArchTag = cutlass::arch::Sm120;
using OperatorClass = cutlass::arch::OpClassBlockScaledTensorOp; // Operator class tag
using KernelMainloopPolicy = cutlass::gemm::KernelTmaWarpSpecializedCooperative;
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
ArchTag, OperatorClass,
MmaTileShape, ClusterShape,
cutlass::epilogue::collective::EpilogueTileAuto,
ElementAccumulator, ElementAccumulator,
ElementC, LayoutCTag, AlignmentC,
ElementD, LayoutDTag, AlignmentD,
cutlass::epilogue::collective::EpilogueScheduleAuto // Epilogue schedule policy
>::CollectiveOp;
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
ArchTag, OperatorClass,
ElementA, LayoutATag, AlignmentA,
ElementB, LayoutBTag, AlignmentB,
ElementAccumulator,
MmaTileShape, ClusterShape,
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
KernelMainloopPolicy>::CollectiveOp;
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
Shape<int, int, int, int>, // Indicates ProblemShape
CollectiveMainloop,
CollectiveEpilogue,
void>;
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
// Reference device GEMM implementation type
using StrideA = typename Gemm::GemmKernel::StrideA;
using StrideB = typename Gemm::GemmKernel::StrideB;
using StrideC = typename Gemm::GemmKernel::StrideC;
using StrideD = typename Gemm::GemmKernel::StrideD;
using LayoutSFA = typename Gemm::GemmKernel::CollectiveMainloop::LayoutSFA; // Scale Factor tensors have an interleaved layout. Bring Layout instead of stride.
using LayoutSFB = typename Gemm::GemmKernel::CollectiveMainloop::LayoutSFB; // Scale Factor tensors have an interleaved layout. Bring Layout instead of stride.
using Sm1xxBlkScaledConfig = typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig;
torch::checkAllContiguous("gemm_fp4fp4_accum_fp32_out_bf16", {{A, "A", 0}, {B, "B", 1}, {A_sf, "A_sf", 2}, {B_sf, "B_sf", 3}, {alpha, "alpha", 4}});
torch::checkDeviceType("gemm_fp4fp4_accum_fp32_out_bf16", {A, B, A_sf, B_sf, alpha}, at::DeviceType::CUDA);
torch::checkAllSameGPU("gemm_fp4fp4_accum_fp32_out_bf16", {{A, "A", 0}, {B, "B", 1}, {A_sf, "A_sf", 2}, {B_sf, "B_sf", 3}, {alpha, "alpha", 4}});
check_block_scale_factor_type<ElementA>(A_sf, "A_sf");
check_block_scale_factor_type<ElementB>(B_sf, "B_sf");
auto [M, N, K] = check_and_get_fp4_matmul_dims<ElementA, LayoutATag, ElementB, LayoutBTag>(A, B, A_sf, B_sf);
auto D = torch::empty({M, N}, torch::dtype(element_traits<ElementD>::scalar_type).device(A.device()));
Gemm gemm;
// Create stride and layout information for the packed tensors
// For packed NVFP4 tensors, we need to use the appropriate stride and layout
StrideA stride_A = cutlass::make_cute_packed_stride(StrideA{}, {M, K, 1});
StrideB stride_B = cutlass::make_cute_packed_stride(StrideB{}, {N, K, 1});
StrideC stride_C = cutlass::make_cute_packed_stride(StrideC{}, {M, N, 1});
StrideD stride_D = cutlass::make_cute_packed_stride(StrideD{}, {M, N, 1});
// Create scale factor layouts
LayoutSFA layout_SFA = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(make_shape(M, N, K, 1));
LayoutSFB layout_SFB = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(make_shape(M, N, K, 1));
typename Gemm::Arguments args{
cutlass::gemm::GemmUniversalMode::kGemm,
{M, N, K, 1},
{// Mainloop arguments
static_cast<typename ElementA::DataType const *>(A.data_ptr()), stride_A,
static_cast<typename ElementB::DataType const *>(B.data_ptr()), stride_B,
static_cast<typename ElementA::ScaleFactorType const *>(A_sf.data_ptr()), layout_SFA,
static_cast<typename ElementB::ScaleFactorType const *>(B_sf.data_ptr()), layout_SFB},
{// Epilogue arguments
{1.0f, 0.0f},
nullptr,
stride_C,
static_cast<ElementD *>(D.data_ptr()),
stride_D}};
args.epilogue.thread.alpha_ptr = static_cast<ElementAccumulator *>(alpha.data_ptr());
// Check if the problem size is supported or not
CUTLASS_CHECK(gemm.can_implement(args));
// Initialize CUTLASS kernel with arguments and workspace pointer
CUTLASS_CHECK(gemm.initialize(args));
CUTLASS_CHECK(gemm.run(args));
return D;
}
// Adapted from example 72b
torch::Tensor gemm_nvfp4nvfp4_accum_fp32_out_bf16_tnt_sm120(torch::Tensor const &A, torch::Tensor const &B, torch::Tensor const &A_sf, torch::Tensor const &B_sf, torch::Tensor const &alpha)
{
at::cuda::CUDAGuard device_guard(A.device());
using ElementA = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
using ElementB = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
using ElementD = cutlass::bfloat16_t;
using MmaTileShape = Shape<_128, _128, _128>;
using ClusterShape = Shape<_1, _1, _1>;
// C/D matrix configuration
using ElementC = void; // Element type for C matrix operand
using LayoutATag = cutlass::layout::RowMajor; // Layout type for A matrix operand
using LayoutBTag = cutlass::layout::ColumnMajor; // Layout type for B matrix operand
using LayoutCTag = cutlass::layout::RowMajor; // Layout type for C matrix operand
using LayoutDTag = cutlass::layout::RowMajor; // Layout type for D matrix operand
constexpr int AlignmentA = 32;
constexpr int AlignmentB = 32;
constexpr int AlignmentC = 1; // Memory access granularity/alignment of C matrix in units of elements (up to 16 bytes)
constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value; // Memory access granularity/alignment of D matrix in units of elements (up to 16 bytes)
// Kernel functional config
using ElementAccumulator = float; // Element type for internal accumulation
using ArchTag = cutlass::arch::Sm120;
using OperatorClass = cutlass::arch::OpClassBlockScaledTensorOp; // Operator class tag
using KernelMainloopPolicy = cutlass::gemm::KernelTmaWarpSpecializedCooperative;
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
ArchTag, OperatorClass,
MmaTileShape, ClusterShape,
cutlass::epilogue::collective::EpilogueTileAuto,
ElementAccumulator, ElementAccumulator,
ElementC, LayoutCTag, AlignmentC,
ElementD, LayoutDTag, AlignmentD,
cutlass::epilogue::collective::EpilogueScheduleAuto // Epilogue schedule policy
>::CollectiveOp;
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
ArchTag, OperatorClass,
ElementA, LayoutATag, AlignmentA,
ElementB, LayoutBTag, AlignmentB,
ElementAccumulator,
MmaTileShape, ClusterShape,
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
KernelMainloopPolicy>::CollectiveOp;
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
Shape<int, int, int, int>, // Indicates ProblemShape
CollectiveMainloop,
CollectiveEpilogue,
void>;
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
// Reference device GEMM implementation type
using StrideA = typename Gemm::GemmKernel::StrideA;
using StrideB = typename Gemm::GemmKernel::StrideB;
using StrideC = typename Gemm::GemmKernel::StrideC;
using StrideD = typename Gemm::GemmKernel::StrideD;
using LayoutSFA = typename Gemm::GemmKernel::CollectiveMainloop::LayoutSFA; // Scale Factor tensors have an interleaved layout. Bring Layout instead of stride.
using LayoutSFB = typename Gemm::GemmKernel::CollectiveMainloop::LayoutSFB; // Scale Factor tensors have an interleaved layout. Bring Layout instead of stride.
using Sm1xxBlkScaledConfig = typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig;
torch::checkAllContiguous("gemm_fp4fp4_accum_fp32_out_bf16", {{A, "A", 0}, {B, "B", 1}, {A_sf, "A_sf", 2}, {B_sf, "B_sf", 3}, {alpha, "alpha", 4}});
torch::checkDeviceType("gemm_fp4fp4_accum_fp32_out_bf16", {A, B, A_sf, B_sf, alpha}, at::DeviceType::CUDA);
torch::checkAllSameGPU("gemm_fp4fp4_accum_fp32_out_bf16", {{A, "A", 0}, {B, "B", 1}, {A_sf, "A_sf", 2}, {B_sf, "B_sf", 3}, {alpha, "alpha", 4}});
check_block_scale_factor_type<ElementA>(A_sf, "A_sf");
check_block_scale_factor_type<ElementB>(B_sf, "B_sf");
auto [M, N, K] = check_and_get_fp4_matmul_dims<ElementA, LayoutATag, ElementB, LayoutBTag>(A, B, A_sf, B_sf);
auto D = torch::empty({M, N}, torch::dtype(element_traits<ElementD>::scalar_type).device(A.device()));
Gemm gemm;
// Create stride and layout information for the packed tensors
// For packed NVFP4 tensors, we need to use the appropriate stride and layout
StrideA stride_A = cutlass::make_cute_packed_stride(StrideA{}, {M, K, 1});
StrideB stride_B = cutlass::make_cute_packed_stride(StrideB{}, {N, K, 1});
StrideC stride_C = cutlass::make_cute_packed_stride(StrideC{}, {M, N, 1});
StrideD stride_D = cutlass::make_cute_packed_stride(StrideD{}, {M, N, 1});
// Create scale factor layouts
LayoutSFA layout_SFA = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(make_shape(M, N, K, 1));
LayoutSFB layout_SFB = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(make_shape(M, N, K, 1));
typename Gemm::Arguments args{
cutlass::gemm::GemmUniversalMode::kGemm,
{M, N, K, 1},
{// Mainloop arguments
static_cast<typename ElementA::DataType const *>(A.data_ptr()), stride_A,
static_cast<typename ElementB::DataType const *>(B.data_ptr()), stride_B,
static_cast<typename ElementA::ScaleFactorType const *>(A_sf.data_ptr()), layout_SFA,
static_cast<typename ElementB::ScaleFactorType const *>(B_sf.data_ptr()), layout_SFB},
{// Epilogue arguments
{1.0f, 0.0f},
nullptr,
stride_C,
static_cast<ElementD *>(D.data_ptr()),
stride_D}};
args.epilogue.thread.alpha_ptr = static_cast<ElementAccumulator *>(alpha.data_ptr());
// Check if the problem size is supported or not
CUTLASS_CHECK(gemm.can_implement(args));
// Initialize CUTLASS kernel with arguments and workspace pointer
CUTLASS_CHECK(gemm.initialize(args));
CUTLASS_CHECK(gemm.run(args));
return D;
}
// Adapted from example 72b
torch::Tensor gemm_nvfp4nvfp4_accum_fp32_out_fp16_tnt_sm120(torch::Tensor const &A, torch::Tensor const &B, torch::Tensor const &A_sf, torch::Tensor const &B_sf, torch::Tensor const &alpha)
{
at::cuda::CUDAGuard device_guard(A.device());
using ElementA = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
using ElementB = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
using ElementD = cutlass::half_t;
using MmaTileShape = Shape<_128, _128, _128>;
using ClusterShape = Shape<_1, _1, _1>;
// C/D matrix configuration
using ElementC = void; // Element type for C matrix operand
using LayoutATag = cutlass::layout::RowMajor; // Layout type for A matrix operand
using LayoutBTag = cutlass::layout::ColumnMajor; // Layout type for B matrix operand
using LayoutCTag = cutlass::layout::RowMajor; // Layout type for C matrix operand
using LayoutDTag = cutlass::layout::RowMajor; // Layout type for D matrix operand
constexpr int AlignmentA = 32;
constexpr int AlignmentB = 32;
constexpr int AlignmentC = 1; // Memory access granularity/alignment of C matrix in units of elements (up to 16 bytes)
constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value; // Memory access granularity/alignment of D matrix in units of elements (up to 16 bytes)
// Kernel functional config
using ElementAccumulator = float; // Element type for internal accumulation
using ArchTag = cutlass::arch::Sm120;
using OperatorClass = cutlass::arch::OpClassBlockScaledTensorOp; // Operator class tag
using KernelMainloopPolicy = cutlass::gemm::KernelTmaWarpSpecializedCooperative;
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
ArchTag, OperatorClass,
MmaTileShape, ClusterShape,
cutlass::epilogue::collective::EpilogueTileAuto,
ElementAccumulator, ElementAccumulator,
ElementC, LayoutCTag, AlignmentC,
ElementD, LayoutDTag, AlignmentD,
cutlass::epilogue::collective::EpilogueScheduleAuto // Epilogue schedule policy
>::CollectiveOp;
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
ArchTag, OperatorClass,
ElementA, LayoutATag, AlignmentA,
ElementB, LayoutBTag, AlignmentB,
ElementAccumulator,
MmaTileShape, ClusterShape,
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
KernelMainloopPolicy>::CollectiveOp;
using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
Shape<int, int, int, int>, // Indicates ProblemShape
CollectiveMainloop,
CollectiveEpilogue,
void>;
using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
// Reference device GEMM implementation type
using StrideA = typename Gemm::GemmKernel::StrideA;
using StrideB = typename Gemm::GemmKernel::StrideB;
using StrideC = typename Gemm::GemmKernel::StrideC;
using StrideD = typename Gemm::GemmKernel::StrideD;
using LayoutSFA = typename Gemm::GemmKernel::CollectiveMainloop::LayoutSFA; // Scale Factor tensors have an interleaved layout. Bring Layout instead of stride.
using LayoutSFB = typename Gemm::GemmKernel::CollectiveMainloop::LayoutSFB; // Scale Factor tensors have an interleaved layout. Bring Layout instead of stride.
using Sm1xxBlkScaledConfig = typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig;
torch::checkAllContiguous("gemm_fp4fp4_accum_fp32_out_bf16", {{A, "A", 0}, {B, "B", 1}, {A_sf, "A_sf", 2}, {B_sf, "B_sf", 3}, {alpha, "alpha", 4}});
torch::checkDeviceType("gemm_fp4fp4_accum_fp32_out_bf16", {A, B, A_sf, B_sf, alpha}, at::DeviceType::CUDA);
torch::checkAllSameGPU("gemm_fp4fp4_accum_fp32_out_bf16", {{A, "A", 0}, {B, "B", 1}, {A_sf, "A_sf", 2}, {B_sf, "B_sf", 3}, {alpha, "alpha", 4}});
check_block_scale_factor_type<ElementA>(A_sf, "A_sf");
check_block_scale_factor_type<ElementB>(B_sf, "B_sf");
auto [M, N, K] = check_and_get_fp4_matmul_dims<ElementA, LayoutATag, ElementB, LayoutBTag>(A, B, A_sf, B_sf);
auto D = torch::empty({M, N}, torch::dtype(element_traits<ElementD>::scalar_type).device(A.device()));
Gemm gemm;
// Create stride and layout information for the packed tensors
// For packed NVFP4 tensors, we need to use the appropriate stride and layout
StrideA stride_A = cutlass::make_cute_packed_stride(StrideA{}, {M, K, 1});
StrideB stride_B = cutlass::make_cute_packed_stride(StrideB{}, {N, K, 1});
StrideC stride_C = cutlass::make_cute_packed_stride(StrideC{}, {M, N, 1});
StrideD stride_D = cutlass::make_cute_packed_stride(StrideD{}, {M, N, 1});
// Create scale factor layouts
LayoutSFA layout_SFA = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(make_shape(M, N, K, 1));
LayoutSFB layout_SFB = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(make_shape(M, N, K, 1));
typename Gemm::Arguments args{
cutlass::gemm::GemmUniversalMode::kGemm,
{M, N, K, 1},
{// Mainloop arguments
static_cast<typename ElementA::DataType const *>(A.data_ptr()), stride_A,
static_cast<typename ElementB::DataType const *>(B.data_ptr()), stride_B,
static_cast<typename ElementA::ScaleFactorType const *>(A_sf.data_ptr()), layout_SFA,
static_cast<typename ElementB::ScaleFactorType const *>(B_sf.data_ptr()), layout_SFB},
{// Epilogue arguments
{1.0f, 0.0f},
nullptr,
stride_C,
static_cast<ElementD *>(D.data_ptr()),
stride_D}};
args.epilogue.thread.alpha_ptr = static_cast<ElementAccumulator *>(alpha.data_ptr());
// Check if the problem size is supported or not
CUTLASS_CHECK(gemm.can_implement(args));
// Initialize CUTLASS kernel with arguments and workspace pointer
CUTLASS_CHECK(gemm.initialize(args));
CUTLASS_CHECK(gemm.run(args));
return D;
}
TORCH_LIBRARY_IMPL(fouroversix, CUDA, m)
{
// See the note at the top of the file for more information about why these
// functions are not templated and not included in fp4_gemm.cu.
/* MXFP4 */
m.impl("gemm_mxfp4mxfp4_accum_fp32_out_bf16_tnt_sm120",
&gemm_mxfp4mxfp4_accum_fp32_out_bf16_tnt_sm120);
/* NVFP4 */
m.impl("gemm_nvfp4nvfp4_accum_fp32_out_bf16_tnt_sm120",
&gemm_nvfp4nvfp4_accum_fp32_out_bf16_tnt_sm120);
m.impl("gemm_nvfp4nvfp4_accum_fp32_out_fp16_tnt_sm120",
&gemm_nvfp4nvfp4_accum_fp32_out_fp16_tnt_sm120);
}
}
fouroversix/src/fouroversix/csrc/include/element_traits.hpp
+171
View File
@@ -0,0 +1,171 @@
#include "cutlass/bfloat16.h"
#include "cutlass/float_subbyte.h"
#include "cutlass/half.h"
#include "cutlass/layout/matrix.h"
#include <ATen/ATen.h>
#include <type_traits>
namespace fouroversix
{
template <typename T>
struct element_traits;
template <>
struct element_traits<cutlass::mx_float4_t<cutlass::float_e2m1_t>>
{
static constexpr size_t block_size = 32;
static constexpr size_t packing_factor = 2;
static constexpr size_t packed_bytes = 1;
static constexpr at::ScalarType block_scale_factor_type = at::kFloat8_e8m0fnu;
static constexpr at::ScalarType scalar_type = at::kByte;
};
template <>
struct element_traits<cutlass::mx_float6_t<cutlass::float_e2m3_t>>
{
static constexpr size_t block_size = 32;
static constexpr size_t packing_factor = 4;
static constexpr size_t packed_bytes = 3;
static constexpr at::ScalarType block_scale_factor_type = at::kFloat8_e8m0fnu;
static constexpr at::ScalarType scalar_type = at::kByte;
};
template <>
struct element_traits<cutlass::mx_float8_t<cutlass::float_e5m2_t>>
{
static constexpr size_t block_size = 32;
static constexpr size_t packing_factor = 1;
static constexpr size_t packed_bytes = 1;
static constexpr at::ScalarType block_scale_factor_type = at::kFloat8_e8m0fnu;
static constexpr at::ScalarType scalar_type = at::kByte;
};
template <>
struct element_traits<cutlass::nv_float4_t<cutlass::float_e2m1_t>>
{
static constexpr size_t block_size = 16;
static constexpr size_t packing_factor = 2;
static constexpr size_t packed_bytes = 1;
static constexpr at::ScalarType block_scale_factor_type = at::kFloat8_e4m3fn;
static constexpr at::ScalarType scalar_type = at::kByte;
};
template <>
struct element_traits<cutlass::bfloat16_t>
{
static constexpr at::ScalarType scalar_type = at::kBFloat16;
};
template <>
struct element_traits<cutlass::half_t>
{
static constexpr at::ScalarType scalar_type = at::kHalf;
};
template <>
struct element_traits<float>
{
static constexpr at::ScalarType scalar_type = at::kFloat;
};
template <typename Element>
void check_block_scale_factor_type(const at::Tensor &t, const char *name)
{
TORCH_CHECK(
t.scalar_type() == element_traits<Element>::block_scale_factor_type,
name, " must be ", at::toString(element_traits<Element>::block_scale_factor_type));
}
template <typename ElementA, typename LayoutATag, typename ElementB, typename LayoutBTag>
std::tuple<int, int, int>
check_and_get_bf16_matmul_dims(const at::Tensor &A, const at::Tensor &B)
{
TORCH_CHECK(A.scalar_type() == at::kBFloat16, "A must be bfloat16");
TORCH_CHECK(B.scalar_type() == at::kBFloat16, "B must be bfloat16");
TORCH_CHECK(A.dim() == 2, "A must be 2D");
TORCH_CHECK(B.dim() == 2, "B must be 2D");
// Unpack sizes
int a_rows = A.size(0), a_cols = A.size(1);
int b_rows = B.size(0), b_cols = B.size(1);
// Layout-based interpretation
int M, K_A;
if constexpr (std::is_same_v<LayoutATag, cutlass::layout::RowMajor>)
{
M = a_rows;
K_A = a_cols;
}
else if constexpr (std::is_same_v<LayoutATag, cutlass::layout::ColumnMajor>)
{
M = a_cols;
K_A = a_rows;
}
int N, K_B;
if constexpr (std::is_same_v<LayoutBTag, cutlass::layout::RowMajor>)
{
K_B = b_rows;
N = b_cols;
}
else if constexpr (std::is_same_v<LayoutBTag, cutlass::layout::ColumnMajor>)
{
K_B = b_cols;
N = b_rows;
}
TORCH_CHECK(K_A == K_B, "Inner dims mismatch: ", K_A, " vs ", K_B);
return {M, N, K_A};
}
template <typename ElementA, typename LayoutATag, typename ElementB, typename LayoutBTag>
std::tuple<int, int, int>
check_and_get_fp4_matmul_dims(const at::Tensor &A, const at::Tensor &B,
const at::Tensor &A_sf, const at::Tensor &B_sf)
{
TORCH_CHECK(A.scalar_type() == at::kByte, "A must be uint8");
TORCH_CHECK(B.scalar_type() == at::kByte, "B must be uint8");
TORCH_CHECK(A.dim() == 2, "A must be 2D");
TORCH_CHECK(B.dim() == 2, "B must be 2D");
TORCH_CHECK(A.size(1) >= 16, "A K-dim must be >= 16");
TORCH_CHECK(B.size(1) >= 16, "B K-dim must be >= 16");
// Unpack 4-bit logical sizes
int a_rows = A.size(0), a_cols = A.size(1) * element_traits<ElementA>::packing_factor / element_traits<ElementA>::packed_bytes;
int b_rows = B.size(0), b_cols = B.size(1) * element_traits<ElementB>::packing_factor / element_traits<ElementB>::packed_bytes;
// Layout-based interpretation
int M, K_A;
if constexpr (std::is_same_v<LayoutATag, cutlass::layout::RowMajor>)
{
M = a_rows;
K_A = a_cols;
}
else if constexpr (std::is_same_v<LayoutATag, cutlass::layout::ColumnMajor>)
{
M = a_cols;
K_A = a_rows;
}
int N, K_B;
if constexpr (std::is_same_v<LayoutBTag, cutlass::layout::RowMajor>)
{
K_B = b_rows;
N = b_cols;
}
else if constexpr (std::is_same_v<LayoutBTag, cutlass::layout::ColumnMajor>)
{
K_B = b_cols;
N = b_rows;
}
TORCH_CHECK(K_A == K_B, "Inner dims mismatch: ", K_A, " vs ", K_B);
// Scale factor checks
TORCH_CHECK(A_sf.numel() * element_traits<ElementA>::block_size == size_t(a_rows) * size_t(a_cols), "A_sf size mismatch");
TORCH_CHECK(B_sf.numel() * element_traits<ElementB>::block_size == size_t(b_rows) * size_t(b_cols), "B_sf size mismatch");
return {M, N, K_A};
}
}
fouroversix/src/fouroversix/csrc/include/fp4_quant.h
+57
View File
@@ -0,0 +1,57 @@
/******************************************************************************
* Copyright (c) 2025, FourOverSix Team.
******************************************************************************/
#pragma once
#include <torch/extension.h>
namespace fouroversix
{
enum AdaptiveBlockScalingRuleType
{
STATIC_6 = 0,
STATIC_4 = 1,
MAE_4o6 = 2,
MSE_4o6 = 3,
ABS_MAX_4o6 = 4,
};
struct FP4_quant_params
{
using index_t = int64_t;
void *__restrict__ x_ptr;
void *__restrict__ x_rht_ptr;
void *__restrict__ x_e2m1_ptr;
void *__restrict__ x_sf_ptr;
void *__restrict__ x_sft_ptr;
void *__restrict__ amax_ptr;
int x_row_stride;
int x_col_stride;
int x_rht_row_stride;
int x_rht_col_stride;
int x_e2m1_row_stride;
int x_e2m1_col_stride;
int x_sf_row_stride;
int x_sf_col_stride;
int x_sft_row_stride;
int x_sft_col_stride;
// The dimensions.
int M, N, M_rounded, N_rounded, M_sf, N_sf;
bool is_bf16;
bool is_nvfp4;
bool is_rtn;
bool is_rht;
bool is_4o6;
bool is_2d;
bool is_transpose;
int selection_rule; // 0: static_6, 1: static_4, 2: 4o6_mae, 3: 4o6_mse
int rbits;
};
template <typename T, bool Is_nvfp4, bool Is_rht, bool Is_transpose>
void run_fp4_quant_(FP4_quant_params &params, cudaStream_t stream);
} // namespace fouroversix
fouroversix/src/fouroversix/csrc/include/fp4_quant_kernel.h
+505
View File
@@ -0,0 +1,505 @@
/******************************************************************************
* Copyright (c) 2024, Tri Dao.
* Adapted by Junxian Guo from https://github.com/Dao-AILab/flash-attention/blob/main/csrc/flash_attn/src/flash_fwd_kernel.h
* Copyright (c) 2025, FourOverSix Team.
******************************************************************************/
#pragma once
// #include "philox_unpack.cuh" // For at::cuda::philox::unpack
#include <cute/tensor.hpp>
#include <type_traits>
#include <cutlass/cutlass.h>
#include <cutlass/array.h>
#include <cutlass/numeric_types.h>
#include "kernel_traits.h"
#include "utils.h"
#include "hadamard_transform.h"
// #include "softmax.h"
// #include "mask.h"
// #include "dropout.h"
// #include "rotary.h"
namespace fouroversix
{
using namespace cute;
template <typename Kernel_traits, bool Is_nvfp4, bool Is_rht, bool Is_2d, bool Is_transpose, bool Is_rtn, int kSelectionRule, typename Params>
inline __device__ void compute_fp4_quant_prologue_block(const Params &params, const int m_block, const int n_block)
{
// Type aliases
using Element = typename Kernel_traits::Element;
using ScaleFactor = typename Kernel_traits::ScaleFactor;
using index_t = typename Kernel_traits::index_t;
// Compile-time constants
constexpr int kGroupN = Kernel_traits::kGroupN;
constexpr int kBlockM = Kernel_traits::kBlockM;
constexpr int kBlockN = Kernel_traits::kBlockN;
constexpr int kNWarps = Kernel_traits::kNWarps;
constexpr int kNumGroupsInRow = Kernel_traits::kNumGroupsInRow;
constexpr int kNumGroupsInCol = Kernel_traits::kNumGroupsInCol;
constexpr float E4M3_MAX_VALUE = Kernel_traits::E4M3_MAX_VALUE;
constexpr float E2M1_MAX_VALUE = Kernel_traits::E2M1_MAX_VALUE;
constexpr AdaptiveBlockScalingRuleType kRule = static_cast<AdaptiveBlockScalingRuleType>(kSelectionRule);
constexpr bool Is_4o6 = kRule == AdaptiveBlockScalingRuleType::MAE_4o6 ||
kRule == AdaptiveBlockScalingRuleType::MSE_4o6 ||
kRule == AdaptiveBlockScalingRuleType::ABS_MAX_4o6;
using VecTypeX = cutlass::Array<Element, kGroupN>;
using VecTypeXFloat = cutlass::Array<float, kGroupN>;
using VecTypeSFT = cutlass::Array<float, 4>;
constexpr int kVecSizeSFT = 4;
// Shared memory
extern __shared__ char smem[];
// Runtime variables
const int tidx = threadIdx.x;
const int num_groups = kNumGroupsInRow * kBlockM;
float *amax_ptr = reinterpret_cast<float *>(params.amax_ptr);
// -------------------------------------------------------------------------
// Tensor Definitions
// -------------------------------------------------------------------------
// Input X (Global Memory)
Tensor mX = make_tensor(make_gmem_ptr(reinterpret_cast<Element *>(params.x_ptr)),
make_shape(params.M, params.N),
make_stride(params.x_row_stride, _1{}));
Tensor gX = local_tile(mX(_, _), Shape<Int<kBlockM>, Int<kBlockN>>{},
make_coord(m_block, n_block));
Tensor mXRHT = make_tensor(make_gmem_ptr(reinterpret_cast<Element *>(params.x_rht_ptr)),
make_shape(params.M_rounded, params.N_rounded),
make_stride(params.x_rht_row_stride, _1{}));
Tensor gXRHT = local_tile(mXRHT(_, _), Shape<Int<kBlockM>, Int<kBlockN>>{},
make_coord(m_block, n_block));
// Scale Factor Temp SFT (Global Memory)
Tensor mSFT = make_tensor(make_gmem_ptr(reinterpret_cast<float *>(params.x_sft_ptr)),
make_shape(params.M, params.N_rounded / kGroupN),
make_stride(params.x_sft_row_stride, _1{}));
Tensor gSFT = local_tile(mSFT(_, _), Shape<Int<kBlockM>, Int<kBlockN / kGroupN>>{},
make_coord(m_block, n_block));
// Shared Memory Tensors
Tensor sX = make_tensor(make_smem_ptr(reinterpret_cast<Element *>(smem)),
typename Kernel_traits::SmemLayoutX{});
// SFT in Shared Memory (placed after X)
Tensor sSFT = make_tensor(make_smem_ptr(reinterpret_cast<float *>(reinterpret_cast<char *>(sX.data().get()) + sizeof(Element) * size(sX))),
typename Kernel_traits::SmemLayoutSFT{});
// -------------------------------------------------------------------------
// Data Loading (X -> Shared)
// -------------------------------------------------------------------------
typename Kernel_traits::GmemTiledCopyX gmem_tiled_copy_X;
auto gmem_thr_copy_X = gmem_tiled_copy_X.get_thread_slice(tidx);
Tensor tXgX = gmem_thr_copy_X.partition_S(gX);
Tensor tXsX = gmem_thr_copy_X.partition_D(sX);
// Construct predicates for bounds checking
Tensor cX = make_identity_tensor(make_shape(size<0>(sX), size<1>(sX)));
Tensor tXcX = gmem_thr_copy_X.partition_S(cX);
Tensor tXpX = make_tensor<bool>(make_shape(size<2>(tXcX)));
for (int i = 0; i < size(tXpX); ++i)
{
tXpX(i) = get<1>(tXcX(0, 0, i)) < params.N - n_block * kBlockN;
}
__syncthreads();
// Async copy from Global to Shared
fouroversix::copy<false, false, true /*Clear_OOB_MN*/, true /*Clear_OOB_K*/>(
gmem_tiled_copy_X, tXgX, tXsX, tXcX, tXpX, params.M - m_block * kBlockM);
cute::cp_async_fence();
fouroversix::cp_async_wait<0>();
__syncthreads();
// -------------------------------------------------------------------------
// Scale Factor Computation
// -------------------------------------------------------------------------
float thr_max = 0.0f;
for (int g_idx = tidx; g_idx < num_groups; g_idx += blockDim.x)
{
const int g_row = g_idx / kNumGroupsInRow;
const int g_col = g_idx % kNumGroupsInRow;
VecTypeXFloat x_vec_float;
for (int i = 0; i < kGroupN; ++i)
{
x_vec_float[i] = static_cast<float>(sX(g_row, g_col * kGroupN + i));
}
if constexpr (Is_rht)
{
hadamard_quant_group<Is_nvfp4, Element>(&x_vec_float[0]);
VecTypeX x_vec;
#pragma unroll
for (int i = 0; i < kGroupN; ++i)
{
// sX(g_row, g_col * kGroupN + i) = static_cast<Element>(x_vec_float[i]);
x_vec[i] = static_cast<Element>(x_vec_float[i]);
}
*reinterpret_cast<VecTypeX *>(&gXRHT(g_row, g_col * kGroupN)) = *reinterpret_cast<VecTypeX *>(&x_vec);
}
// VecTypeX x_vec = *reinterpret_cast<VecTypeX *>(&sX(g_row, g_col * kGroupN));
// Compute max absolute value in group
float sf = 0.0f;
#pragma unroll
for (int i = 0; i < kGroupN; ++i)
{
sf = max(sf, abs(x_vec_float[i]));
}
thr_max = max(thr_max, sf);
sSFT(g_row, g_col) = sf;
}
if constexpr (Is_2d)
{
__syncthreads();
}
if constexpr (Is_2d)
{
MaxOp<float> max_op;
for (int g_idx = tidx; g_idx < num_groups; g_idx += blockDim.x)
{
const int g_row = g_idx % kNumGroupsInCol;
const int g_col = g_idx / kNumGroupsInCol;
float sf = sSFT(g_row, g_col);
float blk_sf = Allreduce<kGroupN>::run(sf, max_op); // kGroupN is 16 or 32
sSFT(g_row, g_col) = blk_sf;
__syncthreads();
}
}
// -------------------------------------------------------------------------
// Normalization Constant Reduction (Block-wide Max)
// -------------------------------------------------------------------------
// Warp-level reduction
MaxOp<float> max_op;
float warp_max = Allreduce<32>::run(thr_max, max_op);
// Block-level reduction via shared memory
float *sRed = reinterpret_cast<float *>(smem);
if (tidx % 32 == 0)
{
sRed[tidx / 32] = warp_max;
}
__syncthreads();
if (tidx == 0)
{
float blk_max = 0.0f;
#pragma unroll
for (int i = 0; i < kNWarps; ++i)
{
blk_max = max(blk_max, sRed[i]);
}
atomicMaxFloat(amax_ptr, blk_max);
}
// -------------------------------------------------------------------------
// Write Back SFT (Shared -> Global)
// -------------------------------------------------------------------------
for (int r_idx = tidx; r_idx < kBlockM; r_idx += blockDim.x)
{
#pragma unroll
for (int i = 0; i < int(kBlockN / kGroupN); i += kVecSizeSFT)
{
*reinterpret_cast<VecTypeSFT *>(&gSFT(r_idx, i)) = *reinterpret_cast<VecTypeSFT *>(&sSFT(r_idx, i));
}
}
}
template <typename Kernel_traits, bool Is_nvfp4, bool Is_rht, bool Is_2d, bool Is_transpose, bool Is_rtn, int kSelectionRule, typename Params>
inline __device__ void compute_fp4_quant_prologue(const Params &params)
{
// TODO: Implement the fp4 quant kernel
const int m_block = blockIdx.x;
// The block index for the batch.
const int n_block = blockIdx.y;
fouroversix::compute_fp4_quant_prologue_block<Kernel_traits, Is_nvfp4, Is_rht, Is_2d, Is_transpose, Is_rtn, kSelectionRule>(params, m_block, n_block);
}
template <typename Kernel_traits, bool Is_nvfp4, bool Is_rht, bool Is_2d, bool Is_transpose, bool Is_rtn, int kSelectionRule, typename Params>
inline __device__ void compute_fp4_quant_block(const Params &params, const int m_block, const int n_block)
{
// Type aliases
using Element = typename Kernel_traits::Element;
using ScaleFactor = typename Kernel_traits::ScaleFactor;
using index_t = typename Kernel_traits::index_t;
// Compile-time constants
constexpr int kGroupN = Kernel_traits::kGroupN;
constexpr int kBlockM = Kernel_traits::kBlockM;
constexpr int kBlockN = Kernel_traits::kBlockN;
constexpr int kBlockMSF = Kernel_traits::kBlockMSF;
constexpr int kBlockNSF = Kernel_traits::kBlockNSF;
constexpr int kNWarps = Kernel_traits::kNWarps;
constexpr int kNumGroupsInRow = Kernel_traits::kNumGroupsInRow;
constexpr int kNumGroupsInCol = Kernel_traits::kNumGroupsInCol;
constexpr float E4M3_MAX_VALUE = Kernel_traits::E4M3_MAX_VALUE;
constexpr AdaptiveBlockScalingRuleType kRule = static_cast<AdaptiveBlockScalingRuleType>(kSelectionRule);
constexpr bool Is_4o6 = kRule == AdaptiveBlockScalingRuleType::MAE_4o6 ||
kRule == AdaptiveBlockScalingRuleType::MSE_4o6 ||
kRule == AdaptiveBlockScalingRuleType::ABS_MAX_4o6;
constexpr float E4M3_SCALE_4 = Is_4o6 ? Kernel_traits::E4M3_MAX_FOUROVERSIX : E4M3_MAX_VALUE;
constexpr float E4M3_SCALE_6 = Is_4o6 ? Kernel_traits::E4M3_MAX_FOUROVERSIX : E4M3_MAX_VALUE;
constexpr float E2M1_SCALE_4 = Is_4o6 ? 6.0f : 4.0f;
constexpr float E2M1_SCALE_6 = 6.0f;
constexpr int kSmemBlockInRow = int(kNumGroupsInRow / 4);
constexpr int kSmemBlockInCol = int(kBlockM / 128);
using VecTypeXe2m1 = std::conditional_t<Is_nvfp4, cutlass::Array<uint8_t, 8>, cutlass::Array<uint8_t, 16>>;
using VecTypeSFT = cutlass::Array<float, 4>;
using VecTypeSF = cutlass::Array<ScaleFactor, 16>;
using OutputType = cutlass::Array<cutlass::float_e2m1_t, 8>;
constexpr int kVecSizeXe2m1 = Is_nvfp4 ? 8 : 16;
constexpr int kVecSizeSFT = 4;
constexpr int kVecSizeSF = 16;
// Shared memory
extern __shared__ char smem[];
// Runtime variables
const int tidx = threadIdx.x;
const int num_groups = kNumGroupsInRow * kBlockM;
// JXGuo: assure amax is not zero before calling this kernel
const float amax = *reinterpret_cast<float *>(params.amax_ptr);
if (amax == 0.0f)
{
return;
}
// -------------------------------------------------------------------------
// Tensor Definitions
// -------------------------------------------------------------------------
// Input X (Global Memory)
// void *__restrict__ x_ptr = Is_rht ? params.x_rht_ptr : params.x_ptr;
Tensor mX = Is_rht ? make_tensor(make_gmem_ptr(reinterpret_cast<Element *>(params.x_rht_ptr)),
make_shape(params.M_rounded, params.N_rounded),
make_stride(params.x_rht_row_stride, _1{}))
: make_tensor(make_gmem_ptr(reinterpret_cast<Element *>(params.x_ptr)),
make_shape(params.M, params.N),
make_stride(params.x_row_stride, _1{}));
Tensor gX = local_tile(mX(_, _), Shape<Int<kBlockM>, Int<kBlockN>>{},
make_coord(m_block, n_block));
Tensor mXe2m1 = make_tensor(make_gmem_ptr(reinterpret_cast<uint8_t *>(params.x_e2m1_ptr)),
make_shape(params.M, params.N_rounded / 2),
make_stride(params.x_e2m1_row_stride, _1{}));
Tensor gXe2m1 = local_tile(mXe2m1(_, _), Shape<Int<kBlockM>, Int<kBlockN / 2>>{},
make_coord(m_block, n_block));
// Scale Factor Temp SFT (Global Memory)
Tensor mSFT = make_tensor(make_gmem_ptr(reinterpret_cast<float *>(params.x_sft_ptr)),
make_shape(params.M, params.N_rounded / kGroupN),
make_stride(params.x_sft_row_stride, _1{}));
Tensor gSFT = local_tile(mSFT(_, _), Shape<Int<kBlockM>, Int<kBlockN / kGroupN>>{},
make_coord(m_block, n_block));
Tensor gSF = make_tensor(make_gmem_ptr(reinterpret_cast<ScaleFactor *>(params.x_sf_ptr)),
make_shape(params.M_sf, params.N_sf),
make_stride(params.x_sf_row_stride, _1{}));
// Tensor gSF = local_tile(mSF(_, _), Shape<Int<1>, Int<16>>{},
// make_coord(m_block, n_block));
// Shared Memory Tensors
Tensor sX = make_tensor(make_smem_ptr(reinterpret_cast<Element *>(smem)),
typename Kernel_traits::SmemLayoutX{});
// SFT in Shared Memory (placed after X)
Tensor sSFT = make_tensor(make_smem_ptr(reinterpret_cast<float *>(reinterpret_cast<char *>(sX.data().get()) + sizeof(Element) * size(sX))),
typename Kernel_traits::SmemLayoutSFT{});
Tensor sXe2m1 = make_tensor(make_smem_ptr(reinterpret_cast<uint8_t *>(reinterpret_cast<char *>(sSFT.data().get()) + sizeof(float) * size(sSFT))),
Shape<Int<kBlockM>, Int<kBlockN / 2>>{},
Stride<Int<kBlockN / 2>, _1>{});
Tensor sSF = make_tensor(make_smem_ptr(reinterpret_cast<ScaleFactor *>(reinterpret_cast<char *>(sXe2m1.data().get()) + sizeof(uint8_t) * size(sXe2m1))),
typename Kernel_traits::SmemLayoutSF{});
// -------------------------------------------------------------------------
// Data Loading (X -> Shared)
// -------------------------------------------------------------------------
typename Kernel_traits::GmemTiledCopyX gmem_tiled_copy_X;
auto gmem_thr_copy_X = gmem_tiled_copy_X.get_thread_slice(tidx);
Tensor tXgX = gmem_thr_copy_X.partition_S(gX);
Tensor tXsX = gmem_thr_copy_X.partition_D(sX);
// Construct predicates for bounds checking
Tensor cX = make_identity_tensor(make_shape(size<0>(sX), size<1>(sX)));
Tensor tXcX = gmem_thr_copy_X.partition_S(cX);
Tensor tXpX = make_tensor<bool>(make_shape(size<2>(tXcX)));
for (int i = 0; i < size(tXpX); ++i)
{
tXpX(i) = get<1>(tXcX(0, 0, i)) < params.N - n_block * kBlockN;
}
__syncthreads();
// Async copy from Global to Shared
fouroversix::copy<false, false, true /*Clear_OOB_MN*/, true /*Clear_OOB_K*/>(
gmem_tiled_copy_X, tXgX, tXsX, tXcX, tXpX, params.M - m_block * kBlockM);
cute::cp_async_fence();
// -------------------------------------------------------------------------
// Data Loading (SFT -> Shared)
// -------------------------------------------------------------------------
for (int r_idx = tidx; r_idx < kBlockM; r_idx += blockDim.x)
{
#pragma unroll
for (int i = 0; i < int(kBlockN / kGroupN); i += kVecSizeSFT)
{
*reinterpret_cast<VecTypeSFT *>(&sSFT(r_idx, i)) = *reinterpret_cast<VecTypeSFT *>(&gSFT(r_idx, i));
}
}
fouroversix::cp_async_wait<0>();
__syncthreads();
// -------------------------------------------------------------------------
// Quantization
// -------------------------------------------------------------------------
for (int g_idx = tidx; g_idx < num_groups; g_idx += blockDim.x)
{
const int g_row = g_idx % kNumGroupsInCol;
const int g_col = g_idx / kNumGroupsInCol;
const float g_max = sSFT(g_row, g_col);
const Tensor sGX = make_tensor(make_smem_ptr(sX.data() + g_row * kBlockN + g_col * kGroupN),
Shape<Int<1>, Int<kGroupN>>{},
Stride<Int<kGroupN>, _1>{});
OutputType res[int(kGroupN / 8)];
float encode_scale;
float sf;
if constexpr (Is_4o6)
{
encode_scale = E2M1_SCALE_6 * E4M3_SCALE_6 / amax;
float sf_high_precision = g_max / E2M1_SCALE_6 * encode_scale;
float sf_[2] = {sf_high_precision * 1.5, sf_high_precision};
sf_[0] = static_cast<float>(static_cast<ScaleFactor>(sf_[0]));
sf_[1] = static_cast<float>(static_cast<ScaleFactor>(sf_[1]));
sf = fp4_conversion<Is_nvfp4, Is_2d, true, Is_rtn, kRule>(sGX, amax, sf_, res, params.rbits);
}
else
{
float sf_val = 0.0f;
if constexpr (kRule == AdaptiveBlockScalingRuleType::STATIC_6)
{
encode_scale = E4M3_SCALE_6 * E2M1_SCALE_6 / amax;
sf_val = clamp(g_max / E2M1_SCALE_6 * encode_scale, 0, E4M3_MAX_VALUE);
}
else if constexpr (kRule == AdaptiveBlockScalingRuleType::STATIC_4)
{
encode_scale = E2M1_SCALE_4 * E4M3_SCALE_4 / amax;
sf_val = clamp(g_max / E2M1_SCALE_4 * encode_scale, 0, E4M3_MAX_VALUE);
}
else
{
printf("in fp4_quant_block, kRule = %d, not supported\n", kRule);
assert(false);
}
sf_val = static_cast<float>(static_cast<ScaleFactor>(sf_val));
sf = fp4_conversion<Is_nvfp4, false, false, Is_rtn, kRule>(sGX, amax, &sf_val, res, params.rbits);
}
// Write quantized data
for (int i = 0; i < int(kGroupN / 8); ++i)
{
*reinterpret_cast<OutputType *>(&sXe2m1(g_row, g_col * (kGroupN / 2) + i * 4)) = res[i];
}
// Write scale factor (layout: 128x4 blocks, 32 rows per block)
const int r_in_blk = g_row % 128;
const int c_in_blk = g_col % 4;
const int blk_row = int(g_row / 128);
const int blk_col = int(g_col / 4);
const int sf_row = 32 * (blk_row * kSmemBlockInRow + blk_col) + r_in_blk % 32;
const int sf_col = int(r_in_blk / 32) * 4 + c_in_blk;
sSF(sf_row, sf_col) = static_cast<ScaleFactor>(sf);
__syncthreads();
}
// -------------------------------------------------------------------------
// Write Back Xe2m1 (Shared -> Global)
// -------------------------------------------------------------------------
__syncthreads();
for (int r_idx = tidx; r_idx < kBlockM; r_idx += blockDim.x)
{
#pragma unroll
for (int i = 0; i < int(kBlockN / 2); i += kVecSizeXe2m1)
{
*reinterpret_cast<VecTypeXe2m1 *>(&gXe2m1(r_idx, i)) = *reinterpret_cast<VecTypeXe2m1 *>(&sXe2m1(r_idx, i));
}
}
// -------------------------------------------------------------------------
// Write Back SF (Shared -> Global)
// -------------------------------------------------------------------------
const int gbl_blk_row_stride = int(params.N_rounded / (kGroupN * 4));
const int gbl_blk_col_stride = 1;
const int gbl_blk_idx_base = (m_block * kSmemBlockInCol) * gbl_blk_row_stride + (n_block * kSmemBlockInRow) * gbl_blk_col_stride;
static_assert(kVecSizeSF == kBlockNSF, "kVecSizeSF must be equal to kBlockNSF");
for (int r_idx = tidx; r_idx < kBlockMSF; r_idx += blockDim.x)
{
const int loc_blk_idx = int(r_idx / 32);
const int loc_row = r_idx % 32;
const int loc_blk_row = int(loc_blk_idx / kSmemBlockInRow);
const int loc_blk_col = int(loc_blk_idx % kSmemBlockInRow);
const int gbl_blk_idx = gbl_blk_idx_base + loc_blk_row * gbl_blk_row_stride + loc_blk_col * gbl_blk_col_stride;
const index_t gbl_row = index_t(32) * gbl_blk_idx + loc_row;
*reinterpret_cast<VecTypeSF *>(&gSF(gbl_row, 0)) = *reinterpret_cast<VecTypeSF *>(&sSF(r_idx, 0));
}
}
template <typename Kernel_traits, bool Is_nvfp4, bool Is_rht, bool Is_2d, bool Is_transpose, bool Is_rtn, int kSelectionRule, typename Params>
inline __device__ void compute_fp4_quant(const Params &params)
{
// TODO: Implement the fp4 quant kernel
const int m_block = blockIdx.x;
// The block index for the batch.
const int n_block = blockIdx.y;
fouroversix::compute_fp4_quant_block<Kernel_traits, Is_nvfp4, Is_rht, Is_2d, Is_transpose, Is_rtn, kSelectionRule>(params, m_block, n_block);
}
} // namespace fouroversix
fouroversix/src/fouroversix/csrc/include/fp4_quant_launch_template.h
+147
View File
@@ -0,0 +1,147 @@
/******************************************************************************
* Copyright (c) 2023, Tri Dao.
* Adapted by Junxian Guo from https://github.com/Dao-AILab/flash-attention/blob/main/csrc/flash_attn/src/flash_fwd_launch_template.h
* Copyright (c) 2025, FourOverSix Team.
******************************************************************************/
#pragma once
#include <c10/cuda/CUDAException.h> // For C10_CUDA_CHECK and C10_CUDA_KERNEL_LAUNCH_CHECK
#include "static_switch.h"
#include "hardware_info.h"
#include "fp4_quant.h"
#include "fp4_quant_kernel.h"
namespace fouroversix
{
// Determine if the architecture supports FLASH and define a macro to handle parameter modifiers
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000
#define ARCH_SUPPORTS_FLASH
#define KERNEL_PARAM_MODIFIER __grid_constant__
#else
#define KERNEL_PARAM_MODIFIER
#endif
// Define a macro for unsupported architecture handling to centralize the error message
#define FLASH_UNSUPPORTED_ARCH printf("FATAL: FourOverSix requires building with sm version sm100, but was built for < 10.0!");
// Use a macro to clean up kernel definitions
#define DEFINE_FP4_QUANT_KERNEL(kernelName, ...) \
template <typename Kernel_traits, __VA_ARGS__> \
__global__ void kernelName(KERNEL_PARAM_MODIFIER const FP4_quant_params params)
DEFINE_FP4_QUANT_KERNEL(fp4_quant_prologue_kernel, bool Is_nvfp4, bool Is_rht, bool Is_2d, bool Is_transpose, bool Is_rtn, int kSelectionRule)
{
#if defined(ARCH_SUPPORTS_FLASH)
fouroversix::compute_fp4_quant_prologue<Kernel_traits, Is_nvfp4, Is_rht, Is_2d, Is_transpose, Is_rtn, kSelectionRule>(params);
#else
FLASH_UNSUPPORTED_ARCH
#endif
}
DEFINE_FP4_QUANT_KERNEL(fp4_quant_kernel, bool Is_nvfp4, bool Is_rht, bool Is_2d, bool Is_transpose, bool Is_rtn, int kSelectionRule)
{
#if defined(ARCH_SUPPORTS_FLASH)
fouroversix::compute_fp4_quant<Kernel_traits, Is_nvfp4, Is_rht, Is_2d, Is_transpose, Is_rtn, kSelectionRule>(params);
#else
FLASH_UNSUPPORTED_ARCH
#endif
}
////////////////////////////////////////////////////////////////////////////////////////////////////
template <typename Kernel_traits, bool Is_nvfp4, bool Is_rht, bool Is_transpose>
void launch_fp4_quant_prologue(FP4_quant_params &params, cudaStream_t stream)
{
constexpr size_t smem_size = Kernel_traits::kSmemSize;
const int num_m_block = (params.M + Kernel_traits::kBlockM - 1) / Kernel_traits::kBlockM;
const int num_n_block = (params.N + Kernel_traits::kBlockN - 1) / Kernel_traits::kBlockN;
dim3 grid(num_m_block, num_n_block);
BOOL_SWITCH(params.is_rtn, Is_rtn, [&]
{
BOOL_SWITCH(params.is_2d, Is_2d, [&]
{
SELECTION_RULE_SWITCH(params.selection_rule, kSelectionRule, [&]
{
auto kernel_prologue = &fp4_quant_prologue_kernel<Kernel_traits, Is_nvfp4, Is_rht, Is_2d, Is_transpose, Is_rtn, kSelectionRule>;
if (smem_size >= 48 * 1024) {
C10_CUDA_CHECK(cudaFuncSetAttribute(
kernel_prologue, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));
}
kernel_prologue<<<grid, Kernel_traits::kNThreads, smem_size, stream>>>(params);
C10_CUDA_KERNEL_LAUNCH_CHECK();
});
});
});
}
////////////////////////////////////////////////////////////////////////////////////////////////////
template <typename Kernel_traits, bool Is_nvfp4, bool Is_rht, bool Is_transpose>
void launch_fp4_quant(FP4_quant_params &params, cudaStream_t stream)
{
constexpr size_t smem_size = Kernel_traits::kSmemSize;
const int num_m_block = (params.M + Kernel_traits::kBlockM - 1) / Kernel_traits::kBlockM;
const int num_n_block = (params.N + Kernel_traits::kBlockN - 1) / Kernel_traits::kBlockN;
dim3 grid(num_m_block, num_n_block);
BOOL_SWITCH(params.is_rtn, Is_rtn, [&]
{
BOOL_SWITCH(params.is_2d, Is_2d, [&]
{
SELECTION_RULE_SWITCH(params.selection_rule, kSelectionRule, [&]
{
auto kernel = &fp4_quant_kernel<Kernel_traits, Is_nvfp4, Is_rht, Is_2d, Is_transpose, Is_rtn, kSelectionRule>;
if (smem_size >= 48 * 1024) {
C10_CUDA_CHECK(cudaFuncSetAttribute(
kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));
}
kernel<<<grid, Kernel_traits::kNThreads, smem_size, stream>>>(params);
C10_CUDA_KERNEL_LAUNCH_CHECK();
});
});
});
}
////////////////////////////////////////////////////////////////////////////////////////////////////
// template<int kBlockM_, int kBlockN_, int kNWarps_, bool Is_nvfp4, bool Is_transpose, typename elem_type=cutlass::half_t, typename Base=Base_kernel_traits<elem_type>>
template <typename T, bool Is_transpose>
void run_mxfp4_quant(FP4_quant_params &params, cudaStream_t stream)
{
constexpr bool Is_nvfp4 = false;
constexpr bool Is_rht = false;
launch_fp4_quant_prologue<FP4_quant_kernel_traits<128, 128, 4, Is_nvfp4, Is_transpose, T>, Is_nvfp4, Is_rht, Is_transpose>(params, stream);
launch_fp4_quant<FP4_quant_kernel_traits<128, 128, 4, Is_nvfp4, Is_transpose, T>, Is_nvfp4, Is_rht, Is_transpose>(params, stream);
}
template <typename T, bool Is_transpose>
void run_mxfp4_quant_rht(FP4_quant_params &params, cudaStream_t stream)
{
constexpr bool Is_nvfp4 = false;
constexpr bool Is_rht = true;
launch_fp4_quant_prologue<FP4_quant_kernel_traits<128, 128, 4, Is_nvfp4, Is_transpose, T>, Is_nvfp4, Is_rht, Is_transpose>(params, stream);
launch_fp4_quant<FP4_quant_kernel_traits<128, 128, 4, Is_nvfp4, Is_transpose, T>, Is_nvfp4, Is_rht, Is_transpose>(params, stream);
}
template <typename T, bool Is_transpose>
void run_nvfp4_quant(FP4_quant_params &params, cudaStream_t stream)
{
constexpr bool Is_nvfp4 = true;
constexpr bool Is_rht = false;
launch_fp4_quant_prologue<FP4_quant_kernel_traits<128, 64, 4, Is_nvfp4, Is_transpose, T>, Is_nvfp4, Is_rht, Is_transpose>(params, stream);
launch_fp4_quant<FP4_quant_kernel_traits<128, 64, 4, Is_nvfp4, Is_transpose, T>, Is_nvfp4, Is_rht, Is_transpose>(params, stream);
}
template <typename T, bool Is_transpose>
void run_nvfp4_quant_rht(FP4_quant_params &params, cudaStream_t stream)
{
constexpr bool Is_nvfp4 = true;
constexpr bool Is_rht = true;
launch_fp4_quant_prologue<FP4_quant_kernel_traits<128, 64, 4, Is_nvfp4, Is_transpose, T>, Is_nvfp4, Is_rht, Is_transpose>(params, stream);
launch_fp4_quant<FP4_quant_kernel_traits<128, 64, 4, Is_nvfp4, Is_transpose, T>, Is_nvfp4, Is_rht, Is_transpose>(params, stream);
}
} // namespace fouroversix
fouroversix/src/fouroversix/csrc/include/hadamard_transform.h
+56
View File
@@ -0,0 +1,56 @@
/******************************************************************************
* Copyright (c) 2023, Tri Dao.
* Adapted by Junxian Guo from https://github.com/Dao-AILab/fast-hadamard-transform/blob/master/csrc/code_gen.py
* Copyright (c) 2025, FourOverSix Team.
******************************************************************************/
// This file is auto-generated. See "hadamard_code_gen.py"
#pragma once
namespace fouroversix {
template <typename Element>
__device__ __forceinline__ void hadamard_quant_group_16(float x[16]) {
float out[16];
out[0] = + x[0] + x[1] + x[2] - x[3] + x[4] - x[5] - x[6] - x[7] - x[8] - x[9] - x[10] + x[11] - x[12] + x[13] - x[14] - x[15];
out[1] = + x[0] - x[1] + x[2] + x[3] + x[4] + x[5] - x[6] + x[7] - x[8] + x[9] - x[10] - x[11] - x[12] - x[13] - x[14] + x[15];
out[2] = + x[0] + x[1] - x[2] + x[3] + x[4] - x[5] + x[6] + x[7] - x[8] - x[9] + x[10] - x[11] - x[12] + x[13] + x[14] + x[15];
out[3] = + x[0] - x[1] - x[2] - x[3] + x[4] + x[5] + x[6] - x[7] - x[8] + x[9] + x[10] + x[11] - x[12] - x[13] + x[14] - x[15];
out[4] = + x[0] + x[1] + x[2] - x[3] - x[4] + x[5] + x[6] + x[7] - x[8] - x[9] - x[10] + x[11] + x[12] - x[13] + x[14] + x[15];
out[5] = + x[0] - x[1] + x[2] + x[3] - x[4] - x[5] + x[6] - x[7] - x[8] + x[9] - x[10] - x[11] + x[12] + x[13] + x[14] - x[15];
out[6] = + x[0] + x[1] - x[2] + x[3] - x[4] + x[5] - x[6] - x[7] - x[8] - x[9] + x[10] - x[11] + x[12] - x[13] - x[14] - x[15];
out[7] = + x[0] - x[1] - x[2] - x[3] - x[4] - x[5] - x[6] + x[7] - x[8] + x[9] + x[10] + x[11] + x[12] + x[13] - x[14] + x[15];
out[8] = + x[0] + x[1] + x[2] - x[3] + x[4] - x[5] - x[6] - x[7] + x[8] + x[9] + x[10] - x[11] + x[12] - x[13] + x[14] + x[15];
out[9] = + x[0] - x[1] + x[2] + x[3] + x[4] + x[5] - x[6] + x[7] + x[8] - x[9] + x[10] + x[11] + x[12] + x[13] + x[14] - x[15];
out[10] = + x[0] + x[1] - x[2] + x[3] + x[4] - x[5] + x[6] + x[7] + x[8] + x[9] - x[10] + x[11] + x[12] - x[13] - x[14] - x[15];
out[11] = + x[0] - x[1] - x[2] - x[3] + x[4] + x[5] + x[6] - x[7] + x[8] - x[9] - x[10] - x[11] + x[12] + x[13] - x[14] + x[15];
out[12] = + x[0] + x[1] + x[2] - x[3] - x[4] + x[5] + x[6] + x[7] + x[8] + x[9] + x[10] - x[11] - x[12] + x[13] - x[14] - x[15];
out[13] = + x[0] - x[1] + x[2] + x[3] - x[4] - x[5] + x[6] - x[7] + x[8] - x[9] + x[10] + x[11] - x[12] - x[13] - x[14] + x[15];
out[14] = + x[0] + x[1] - x[2] + x[3] - x[4] + x[5] - x[6] - x[7] + x[8] + x[9] - x[10] + x[11] - x[12] + x[13] + x[14] + x[15];
out[15] = + x[0] - x[1] - x[2] - x[3] - x[4] - x[5] - x[6] + x[7] + x[8] - x[9] - x[10] - x[11] - x[12] - x[13] + x[14] - x[15];
#pragma unroll
for (int i = 0; i < 16; i++) { x[i] = static_cast<float>(static_cast<Element>(out[i] / 4)); }
}
template <typename Element>
__device__ __forceinline__ void hadamard_quant_group_32(float x[32]) {
hadamard_quant_group_16<Element>(x);
hadamard_quant_group_16<Element>(x + 16);
}
template <bool Is_nvfp4, typename Element>
__device__ __forceinline__ void hadamard_quant_group(float* x) {
if constexpr (Is_nvfp4) {
hadamard_quant_group_16<Element>(x);
} else {
hadamard_quant_group_32<Element>(x);
}
}
} // namespace fouroversix
fouroversix/src/fouroversix/csrc/include/hardware_info.h
+41
View File
@@ -0,0 +1,41 @@
/******************************************************************************
* Copyright (c) 2024, Tri Dao.
******************************************************************************/
#pragma once
#include <tuple>
#if !defined(__CUDACC_RTC__)
#include "cuda_runtime.h"
#endif
#define CHECK_CUDA(call) \
do { \
cudaError_t status_ = call; \
if (status_ != cudaSuccess) { \
fprintf(stderr, "CUDA error (%s:%d): %s\n", __FILE__, __LINE__, \
cudaGetErrorString(status_)); \
exit(1); \
} \
} while (0)
inline int get_current_device() {
int device;
CHECK_CUDA(cudaGetDevice(&device));
return device;
}
inline std::tuple<int, int> get_compute_capability(int device) {
int capability_major, capability_minor;
CHECK_CUDA(cudaDeviceGetAttribute(&capability_major, cudaDevAttrComputeCapabilityMajor, device));
CHECK_CUDA(cudaDeviceGetAttribute(&capability_minor, cudaDevAttrComputeCapabilityMinor, device));
return {capability_major, capability_minor};
}
inline int get_num_sm(int device) {
int multiprocessor_count;
CHECK_CUDA(cudaDeviceGetAttribute(&multiprocessor_count, cudaDevAttrMultiProcessorCount, device));
return multiprocessor_count;
}
fouroversix/src/fouroversix/csrc/include/kernel_traits.h
+141
View File
@@ -0,0 +1,141 @@
/******************************************************************************
* Copyright (c) 2024, Tri Dao.
* Adapted by Junxian Guo from https://github.com/Dao-AILab/flash-attention/blob/main/csrc/flash_attn/src/kernel_traits.h
* Copyright (c) 2025, FourOverSix Team.
******************************************************************************/
#pragma once
#include "cute/tensor.hpp"
#include "cutlass/cutlass.h"
#include "cutlass/layout/layout.h"
#include <cutlass/numeric_types.h>
using namespace cute;
template <bool Is_nvfp4, typename elem_type = cutlass::half_t>
struct Base_kernel_traits
{
static constexpr float E2M1_MAX_VALUE = 6.0f;
static constexpr float E4M3_MIN_VALUE = -448.0f;
static constexpr float E4M3_MAX_VALUE = 448.0f;
static constexpr float E4M3_MAX_FOUROVERSIX = 256.0f;
using Element = elem_type;
using ScaleFactor = std::conditional_t<Is_nvfp4, cutlass::float_e4m3_t, uint8_t>;
// using ElementXe2m1Packed = std::conditional_t<Is_nvfp4, uint64_t, uint128_t>;
using NormConst = float;
static constexpr bool Has_cp_async = true;
using index_t = int64_t;
using MMA_Atom_Arch = std::conditional_t<
std::is_same_v<elem_type, cutlass::half_t>,
MMA_Atom<SM80_16x8x16_F32F16F16F32_TN>,
MMA_Atom<SM80_16x8x16_F32BF16BF16F32_TN>>;
using SmemCopyAtom = Copy_Atom<SM75_U32x4_LDSM_N, elem_type>;
using SmemCopyAtomTransposed = Copy_Atom<SM75_U16x8_LDSM_T, elem_type>;
};
// If Share_Q_K_smem is true, that forces Is_Q_in_regs to be true
template <int kBlockM_, int kBlockN_, int kNWarps_, bool Is_nvfp4, bool Is_transpose, typename elem_type = cutlass::half_t, typename Base = Base_kernel_traits<Is_nvfp4, elem_type>>
struct FP4_quant_kernel_traits : public Base
{
static constexpr float E2M1_MAX_VALUE = Base::E2M1_MAX_VALUE;
static constexpr float E2M1_MAX_FOUR = 4;
static constexpr float E4M3_MIN_VALUE = Base::E4M3_MIN_VALUE;
static constexpr float E4M3_MAX_VALUE = Base::E4M3_MAX_VALUE;
static constexpr float E4M3_MAX_FOUROVERSIX = Base::E4M3_MAX_FOUROVERSIX;
using Element = typename Base::Element;
using ScaleFactor = typename Base::ScaleFactor;
using NormConst = typename Base::NormConst;
using index_t = typename Base::index_t;
static constexpr bool Has_cp_async = Base::Has_cp_async;
using SmemCopyAtom = typename Base::SmemCopyAtom;
using SmemCopyAtomTransposed = typename Base::SmemCopyAtomTransposed;
// The number of threads.
static constexpr int kNWarps = kNWarps_;
static constexpr int kNThreads = kNWarps * 32;
static constexpr int kBlockM = kBlockM_;
static constexpr int kBlockN = kBlockN_;
static constexpr int kGroupN = Is_nvfp4 ? 16 : 32; // 16 or 32 elements
static_assert(kBlockM % 128 == 0);
static_assert(kBlockN % (kGroupN * 4) == 0);
static_assert(kBlockN % 64 == 0);
static constexpr int kBlockNSmem = 64; // each cache line is 128 bytes, so we need to align to 64 bytes
static constexpr int kBlockNGmem = kBlockN % 128 == 0 ? 128 : 64;
using TiledMma = TiledMMA<
typename Base::MMA_Atom_Arch,
Layout<Shape<Int<kNWarps>, _1, _1>>, // 4x1x1 or 8x1x1 thread group
Tile<Int<16 * kNWarps>, _16, _16>>;
// static constexpr int kGroupN = Is_nvfp4 ? 16 : 32; // 16 or 32 elements
// static constexpr int kSwizzleM = Is_nvfp4 ? 4 : 5; // 16 or 32 elements
// static constexpr int kSwizzleS = Is_nvfp4 ? 2 : 1; // 4 or 2 elements
// static constexpr int kSwizzleB = Is_nvfp4 ? 2 : 1; // 2 or 1 bits
static constexpr int kNumGroupsInRow = kBlockN / kGroupN;
static_assert(kBlockM % kGroupN == 0, "kBlockM must be a multiple of kGroupN if is 2d");
static constexpr int kNumGroupsInCol = kBlockM; // for 2d scale factor
// static constexpr int kSwizzleM = 3;
// static constexpr int kSwizzleS = 3;
// static constexpr int kSwizzleB = 2;
// using SmemLayoutAtomX = decltype(
// composition(Swizzle<kSwizzleB, kSwizzleM, kSwizzleS>{},
// // This has to be kBlockNSmem, using kHeadDim gives wrong results for d=128
// Layout<Shape<_8, Int<kBlockNSmem>>,
// Stride<Int<kBlockNSmem>, _1>>{}));
// using SmemLayoutX = decltype(tile_to_shape(
// SmemLayoutAtomX{},
// Shape<Int<kBlockM>, Int<kBlockN>>{}));
using SmemLayoutX = Layout<Shape<Int<kBlockM>, Int<kBlockN>>, Stride<Int<kBlockN>, _1>>;
using SmemLayoutXTransposed = decltype(composition(SmemLayoutX{}, make_layout(Shape<Int<kBlockN>, Int<kBlockM>>{}, GenRowMajor{})));
using SmemLayoutXTransposedNoSwizzle = decltype(get_nonswizzle_portion(SmemLayoutXTransposed{}));
using SmemLayoutSFT = Layout<Shape<Int<kBlockM>, Int<kBlockN / kGroupN>>,
Stride<Int<kBlockN / kGroupN>, _1>>;
static constexpr int kBlockMSF = kBlockM / 128 * 32 * int(kBlockN / (kGroupN * 4));
static constexpr int kBlockNSF = 16;
using SmemLayoutSF = Layout<Shape<Int<kBlockMSF>, Int<kBlockNSF>>,
Stride<Int<kBlockNSF>, _1>>;
using SmemLayout = SmemLayoutX;
static constexpr int kSmemXSize = size(SmemLayout{}) * sizeof(Element);
static constexpr int kSmemXe2m1Size = kSmemXSize / 4;
static constexpr int kSmemSFTSize = size(SmemLayoutSFT{}) * sizeof(float);
static constexpr int kSmemSFSize = size(SmemLayoutSF{}) * sizeof(ScaleFactor);
static constexpr int kSmemSize = kSmemXSize + kSmemXe2m1Size + kSmemSFTSize + kSmemSFSize;
static constexpr int kGmemElemsPerLoad = sizeof(cute::uint128_t) / sizeof(Element);
static_assert(kBlockN % kGmemElemsPerLoad == 0, "kBlockN must be a multiple of kGmemElemsPerLoad");
static constexpr int kGmemThreadsPerRow = kBlockNSmem / kGmemElemsPerLoad;
static_assert(kNThreads % kGmemThreadsPerRow == 0, "kNThreads must be a multiple of kGmemThreadsPerRow");
using GmemLayoutAtomX = Layout<Shape<Int<kNThreads / kGmemThreadsPerRow>, Int<kGmemThreadsPerRow>>,
Stride<Int<kGmemThreadsPerRow>, _1>>;
using Gmem_copy_atom_x = Copy_Atom<SM80_CP_ASYNC_CACHEGLOBAL<cute::uint128_t>, Element>;
using GmemTiledCopyX = decltype(make_tiled_copy(Gmem_copy_atom_x{},
GmemLayoutAtomX{},
Layout<Shape<_1, _8>>{}));
using Gmem_copy_atom_sft = Copy_Atom<AutoVectorizingCopyWithAssumedAlignment<64>, float>;
using GmemLayoutAtomSFT = Layout<
Shape<_64, _1>,
Stride<_1, _0>>;
using GmemTiledCopySFT = decltype(make_tiled_copy(Gmem_copy_atom_sft{}, GmemLayoutAtomSFT{}, Layout<Shape<_1, _4>>{}));
};
////////////////////////////////////////////////////////////////////////////////////////////////////
fouroversix/src/fouroversix/csrc/include/static_switch.h
+61
View File
@@ -0,0 +1,61 @@
// Inspired by
// https://github.com/NVIDIA/DALI/blob/main/include/dali/core/static_switch.h
// and https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/Dispatch.h
// Adapted by Junxian Guo from https://github.com/NVIDIA/DALI/blob/main/include/dali/core/static_switch.h
// Copyright (c) 2025, FourOverSix Team.
#pragma once
/// @param COND - a boolean expression to switch by
/// @param CONST_NAME - a name given for the constexpr bool variable.
/// @param ... - code to execute for true and false
///
/// Usage:
/// ```
/// BOOL_SWITCH(flag, BoolConst, [&] {
/// some_function<BoolConst>(...);
/// });
/// ```
#define BOOL_SWITCH(COND, CONST_NAME, ...) \
[&] { \
if (COND) { \
constexpr static bool CONST_NAME = true; \
return __VA_ARGS__(); \
} else { \
constexpr static bool CONST_NAME = false; \
return __VA_ARGS__(); \
} }()
#define FP16_SWITCH(COND, ...) \
[&] { \
if (COND) { \
using fp16_type = cutlass::half_t; \
return __VA_ARGS__(); \
} else { \
using fp16_type = cutlass::bfloat16_t; \
return __VA_ARGS__(); \
} }()
#define SELECTION_RULE_SWITCH(SELECTION_RULE, ...) \
[&] { \
if (SELECTION_RULE == 0) { \
constexpr static int kSelectionRule = 0; \
return __VA_ARGS__(); \
} else if (SELECTION_RULE == 1) { \
constexpr static int kSelectionRule = 1; \
return __VA_ARGS__(); \
} else if (SELECTION_RULE == 2) { \
constexpr static int kSelectionRule = 2; \
return __VA_ARGS__(); \
} else if (SELECTION_RULE == 3) { \
constexpr static int kSelectionRule = 3; \
return __VA_ARGS__(); \
} else if (SELECTION_RULE == 4) { \
constexpr static int kSelectionRule = 4; \
return __VA_ARGS__(); \
} else { \
constexpr static int kSelectionRule = 0; \
return __VA_ARGS__(); \
} }()
fouroversix/src/fouroversix/csrc/include/utils.h
+701
View File
@@ -0,0 +1,701 @@
/******************************************************************************
* Copyright (c) 2023, Tri Dao.
* Adapted by Junxian Guo and Jack Cook from https://github.com/Dao-AILab/flash-attention/blob/main/csrc/flash_attn/src/utils.h
* Copyright (c) 2025, FourOverSix Team.
******************************************************************************/
#pragma once
#include <assert.h>
#include <stdint.h>
#include <stdlib.h>
#include <cuda_fp16.h>
#include <type_traits>
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
#include <cuda_bf16.h>
#endif
#include <cute/tensor.hpp>
#include <cutlass/array.h>
#include <cutlass/cutlass.h>
#include <cutlass/numeric_conversion.h>
#include <cutlass/numeric_types.h>
////////////////////////////////////////////////////////////////////////////////////////////////////
namespace fouroversix
{
////////////////////////////////////////////////////////////////////////////////////////////////////
template <typename T>
__forceinline__ __device__ uint32_t relu2(const uint32_t x);
template <>
__forceinline__ __device__ uint32_t relu2<cutlass::half_t>(const uint32_t x)
{
uint32_t res;
const uint32_t zero = 0u;
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
asm volatile("max.f16x2 %0, %1, %2;\n" : "=r"(res) : "r"(x), "r"(zero));
#else
asm volatile(
"{\n"
"\t .reg .f16x2 sela;\n"
"\t set.gtu.u32.f16x2 sela, %1, %2;\n"
"\t and.b32 %0, sela, %1;\n"
"}\n" : "=r"(res) : "r"(x), "r"(zero));
#endif
return res;
}
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
template <>
__forceinline__ __device__ uint32_t relu2<cutlass::bfloat16_t>(const uint32_t x)
{
uint32_t res;
const uint32_t zero = 0u;
asm volatile("max.bf16x2 %0, %1, %2;\n" : "=r"(res) : "r"(x), "r"(zero));
return res;
}
#endif
////////////////////////////////////////////////////////////////////////////////////////////////////
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
template <typename T>
__forceinline__ __device__ uint32_t convert_relu2(const float2 x);
template <>
__forceinline__ __device__ uint32_t convert_relu2<cutlass::half_t>(const float2 x)
{
uint32_t res;
const uint32_t a = reinterpret_cast<const uint32_t &>(x.x);
const uint32_t b = reinterpret_cast<const uint32_t &>(x.y);
asm volatile("cvt.rn.relu.f16x2.f32 %0, %1, %2;\n" : "=r"(res) : "r"(b), "r"(a));
return res;
}
template <>
__forceinline__ __device__ uint32_t convert_relu2<cutlass::bfloat16_t>(const float2 x)
{
uint32_t res;
const uint32_t a = reinterpret_cast<const uint32_t &>(x.x);
const uint32_t b = reinterpret_cast<const uint32_t &>(x.y);
asm volatile("cvt.rn.relu.bf16x2.f32 %0, %1, %2;\n" : "=r"(res) : "r"(b), "r"(a));
return res;
}
#endif
////////////////////////////////////////////////////////////////////////////////////////////////////
template <typename T>
struct MaxOp
{
__device__ __forceinline__ T operator()(T const &x, T const &y) { return x > y ? x : y; }
};
template <>
struct MaxOp<float>
{
// This is slightly faster
__device__ __forceinline__ float operator()(float const &x, float const &y) { return max(x, y); }
};
////////////////////////////////////////////////////////////////////////////////////////////////////
template <typename T>
struct SumOp
{
__device__ __forceinline__ T operator()(T const &x, T const &y) { return x + y; }
};
////////////////////////////////////////////////////////////////////////////////////////////////////
template <int THREADS>
struct Allreduce
{
static_assert(THREADS == 32 || THREADS == 16 || THREADS == 8 || THREADS == 4);
template <typename T, typename Operator>
static __device__ __forceinline__ T run(T x, Operator &op)
{
constexpr int OFFSET = THREADS / 2;
x = op(x, __shfl_xor_sync(uint32_t(-1), x, OFFSET));
return Allreduce<OFFSET>::run(x, op);
}
};
////////////////////////////////////////////////////////////////////////////////////////////////////
template <>
struct Allreduce<2>
{
template <typename T, typename Operator>
static __device__ __forceinline__ T run(T x, Operator &op)
{
x = op(x, __shfl_xor_sync(uint32_t(-1), x, 1));
return x;
}
};
////////////////////////////////////////////////////////////////////////////////////////////////////
template <bool A_in_regs = false, bool B_in_regs = false, typename Tensor0, typename Tensor1,
typename Tensor2, typename Tensor3, typename Tensor4,
typename TiledMma, typename TiledCopyA, typename TiledCopyB,
typename ThrCopyA, typename ThrCopyB>
__forceinline__ __device__ void gemm(Tensor0 &acc, Tensor1 &tCrA, Tensor2 &tCrB, Tensor3 const &tCsA,
Tensor4 const &tCsB, TiledMma tiled_mma,
TiledCopyA smem_tiled_copy_A, TiledCopyB smem_tiled_copy_B,
ThrCopyA smem_thr_copy_A, ThrCopyB smem_thr_copy_B)
{
CUTE_STATIC_ASSERT_V(size<1>(tCrA) == size<1>(acc)); // MMA_M
CUTE_STATIC_ASSERT_V(size<1>(tCrB) == size<2>(acc)); // MMA_N
CUTE_STATIC_ASSERT_V(size<2>(tCrA) == size<2>(tCrB)); // MMA_K
Tensor tCrA_copy_view = smem_thr_copy_A.retile_D(tCrA);
CUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(tCrA_copy_view)); // M
Tensor tCrB_copy_view = smem_thr_copy_B.retile_D(tCrB);
CUTE_STATIC_ASSERT_V(size<1>(tCsB) == size<1>(tCrB_copy_view)); // N
if (!A_in_regs)
{
cute::copy(smem_tiled_copy_A, tCsA(_, _, _0{}), tCrA_copy_view(_, _, _0{}));
}
if (!B_in_regs)
{
cute::copy(smem_tiled_copy_B, tCsB(_, _, _0{}), tCrB_copy_view(_, _, _0{}));
}
#pragma unroll
for (int i = 0; i < size<2>(tCrA); ++i)
{
if (i < size<2>(tCrA) - 1)
{
if (!A_in_regs)
{
cute::copy(smem_tiled_copy_A, tCsA(_, _, i + 1), tCrA_copy_view(_, _, i + 1));
}
if (!B_in_regs)
{
cute::copy(smem_tiled_copy_B, tCsB(_, _, i + 1), tCrB_copy_view(_, _, i + 1));
}
}
cute::gemm(tiled_mma, tCrA(_, _, i), tCrB(_, _, i), acc);
}
}
////////////////////////////////////////////////////////////////////////////////////////////////////
template <typename Tensor0, typename Tensor1, typename Tensor2, typename Tensor3,
typename TiledMma, typename TiledCopy, typename ThrCopy>
__forceinline__ __device__ void gemm_rs(Tensor0 &acc, Tensor1 &tCrA, Tensor2 &tCrB, Tensor3 const &tCsB,
TiledMma tiled_mma, TiledCopy smem_tiled_copy_B,
ThrCopy smem_thr_copy_B)
{
CUTE_STATIC_ASSERT_V(size<1>(tCrA) == size<1>(acc)); // MMA_M
CUTE_STATIC_ASSERT_V(size<1>(tCrB) == size<2>(acc)); // MMA_N
CUTE_STATIC_ASSERT_V(size<2>(tCrA) == size<2>(tCrB)); // MMA_K
Tensor tCrB_copy_view = smem_thr_copy_B.retile_D(tCrB);
CUTE_STATIC_ASSERT_V(size<1>(tCsB) == size<1>(tCrB_copy_view)); // N
cute::copy(smem_tiled_copy_B, tCsB(_, _, _0{}), tCrB_copy_view(_, _, _0{}));
#pragma unroll
for (int i = 0; i < size<2>(tCrA); ++i)
{
if (i < size<2>(tCrA) - 1)
{
cute::copy(smem_tiled_copy_B, tCsB(_, _, i + 1), tCrB_copy_view(_, _, i + 1));
}
cute::gemm(tiled_mma, tCrA(_, _, i), tCrB(_, _, i), acc);
}
}
////////////////////////////////////////////////////////////////////////////////////////////////////
// Convert acc_layout from (MMA=4, MMA_M, MMA_N) to (nrow=(2, MMA_M), ncol=(2, MMA_N))
template <typename Layout>
__forceinline__ __device__ auto convert_layout_acc_rowcol(Layout acc_layout)
{
static_assert(decltype(size<0>(acc_layout))::value == 4);
static_assert(decltype(rank(acc_layout))::value == 3);
auto l = logical_divide(acc_layout, Shape<_2>{}); // ((2, 2), MMA_M, MMA_N)
return make_layout(make_layout(get<0, 1>(l), get<1>(l)), make_layout(get<0, 0>(l), get<2>(l)));
};
////////////////////////////////////////////////////////////////////////////////////////////////////
// Convert acc_layout from (MMA=4, MMA_M, MMA_N) to ((4, 2), MMA_M, MMA_N / 2)
// if using m16n8k16, or to (4, MMA_M, MMA_N) if using m16n8k8.
template <typename MMA_traits, typename Layout>
__forceinline__ __device__ auto convert_layout_acc_Aregs(Layout acc_layout)
{
using X = Underscore;
static_assert(decltype(size<0>(acc_layout))::value == 4);
static_assert(decltype(rank(acc_layout))::value == 3);
constexpr int mma_shape_K = get<2>(typename MMA_traits::Shape_MNK{});
static_assert(mma_shape_K == 8 || mma_shape_K == 16);
if constexpr (mma_shape_K == 8)
{
return acc_layout;
}
else
{
auto l = logical_divide(acc_layout, Shape<X, X, _2>{}); // (4, MMA_M, (2, MMA_N / 2)))
return make_layout(make_layout(get<0>(l), get<2, 0>(l)), get<1>(l), get<2, 1>(l));
}
};
////////////////////////////////////////////////////////////////////////////////////////////////////
// Convert acc_layout from (MMA=4, MMA_M, MMA_N) to ((4, 2), MMA_M, MMA_N / 2)
template <typename Layout>
__forceinline__ __device__ auto convert_layout_acc_dropout(Layout acc_layout)
{
using X = Underscore;
static_assert(decltype(size<0>(acc_layout))::value == 4);
static_assert(decltype(rank(acc_layout))::value == 3);
auto l = logical_divide(acc_layout, Shape<X, X, _2>{}); // (4, MMA_M, (2, MMA_N / 2)))
return make_layout(make_layout(get<0>(l), get<2, 0>(l)), get<1>(l), get<2, 1>(l));
};
////////////////////////////////////////////////////////////////////////////////////////////////////
template <typename To_type, typename Engine, typename Layout>
__forceinline__ __device__ auto convert_type(Tensor<Engine, Layout> const &tensor)
{
using From_type = typename Engine::value_type;
constexpr int numel = decltype(size(tensor))::value;
cutlass::NumericArrayConverter<To_type, From_type, numel> convert_op;
// HACK: this requires tensor to be "contiguous"
auto frag = convert_op(*reinterpret_cast<const cutlass::Array<From_type, numel> *>(tensor.data()));
return make_tensor(make_rmem_ptr<To_type>(&frag), tensor.layout());
}
////////////////////////////////////////////////////////////////////////////////////////////////////
// Blocks until all but N previous cp.async.commit_group operations have committed.
// This differs from cute::cp_async_wait in that when N = 0 we don't call cp.async.wait_all
// (which is equivalent to commit_group then wait_group 0).
// Instead we just call cp.async.wait_group 0, which is slightly faster.
// https://github.com/NVIDIA/cutlass/blob/master/include/cute/arch/copy_sm80.hpp#L113
template <int N>
CUTE_HOST_DEVICE void cp_async_wait()
{
#if defined(CUTE_ARCH_CP_ASYNC_SM80_ENABLED)
asm volatile("cp.async.wait_group %0;\n" ::"n"(N));
#endif
}
////////////////////////////////////////////////////////////////////////////////////////////////////
template <bool Is_even_MN = true, bool Is_even_K = true, bool Clear_OOB_MN = false, bool Clear_OOB_K = true,
typename TiledCopy, typename Engine0, typename Layout0, typename Engine1, typename Layout1,
typename Engine2, typename Layout2, typename Engine3, typename Layout3>
__forceinline__ __device__ void copy(TiledCopy tiled_copy, Tensor<Engine0, Layout0> const &S,
Tensor<Engine1, Layout1> &D, Tensor<Engine2, Layout2> const &identity_MN,
Tensor<Engine3, Layout3> const &predicate_K, const int max_MN = 0)
{
CUTE_STATIC_ASSERT_V(rank(S) == Int<3>{});
CUTE_STATIC_ASSERT_V(rank(D) == Int<3>{});
CUTE_STATIC_ASSERT_V(size<0>(S) == size<0>(D)); // MMA
CUTE_STATIC_ASSERT_V(size<1>(S) == size<1>(D)); // MMA_M
CUTE_STATIC_ASSERT_V(size<2>(S) == size<2>(D)); // MMA_K
// There's no case where !Clear_OOB_K && Clear_OOB_MN
static_assert(!(Clear_OOB_MN && !Clear_OOB_K));
#pragma unroll
for (int m = 0; m < size<1>(S); ++m)
{
if (Is_even_MN || get<0>(identity_MN(0, m, 0)) < max_MN)
{
#pragma unroll
for (int k = 0; k < size<2>(S); ++k)
{
if (Is_even_K || predicate_K(k))
{
cute::copy(tiled_copy, S(_, m, k), D(_, m, k));
}
else if (Clear_OOB_K)
{
cute::clear(D(_, m, k));
}
}
}
else if (Clear_OOB_MN)
{
cute::clear(D(_, m, _));
}
}
}
static __device__ __forceinline__ float atomicMaxFloat(float *addr, float value)
{
// source: https://stackoverflow.com/a/51549250
return (value >= 0)
? __int_as_float(atomicMax((int *)addr, __float_as_int(value)))
: __uint_as_float(atomicMin((unsigned int *)addr, __float_as_uint(value)));
}
////////////////////////////////////////////////////////////////////////////////////////////////////
template <bool Is_4o6, bool Is_rtn, AdaptiveBlockScalingRuleType kAdaptiveBlockScalingRuleType>
struct Fp4ArrayQuant
{
using InputType = cutlass::Array<float, 8>;
using OutputType = cutlass::Array<cutlass::float_e2m1_t, 8>;
using ScaleFactorType = float;
using ErrorType = float;
__device__ __forceinline__
OutputType
convert(InputType const &x,
const float amax,
const ScaleFactorType sf,
const uint32_t rbits,
// Usage depends on Is_rtn
ErrorType *err /*nullable*/)
{
InputType x_scaled;
constexpr float E2M1_MAX_VALUE = 6.0f;
constexpr float E2M1_MAX_FOUR = 4.0f;
constexpr float E4M3_MAX_VALUE = 448.0f;
constexpr float E4M3_MAX_FOUROVERSIX = 256.0f;
constexpr float e2m1_limit = kAdaptiveBlockScalingRuleType == AdaptiveBlockScalingRuleType::STATIC_4 ? E2M1_MAX_FOUR : E2M1_MAX_VALUE;
constexpr float e4m3_limit = (kAdaptiveBlockScalingRuleType == AdaptiveBlockScalingRuleType::STATIC_6 || kAdaptiveBlockScalingRuleType == AdaptiveBlockScalingRuleType::STATIC_4)
? E4M3_MAX_VALUE
: E4M3_MAX_FOUROVERSIX;
const float encode_scale = e4m3_limit * e2m1_limit / amax;
const float decode_scale = 1.0 / encode_scale;
const float block_scale_inv = fminf(1.0f / (decode_scale * sf), std::numeric_limits<float>::max());
#pragma unroll
for (int i = 0; i < 8; ++i)
{
x_scaled[i] = x[i] * block_scale_inv;
}
unsigned out;
if constexpr (Is_rtn)
{
if constexpr (Is_4o6)
{
unsigned out_dequant_1;
unsigned out_dequant_2;
unsigned out_dequant_3;
unsigned out_dequant_4;
asm volatile(
"{\n"
".reg .b8 byte0, byte1, byte2, byte3;\n"
"cvt.rn.satfinite.e2m1x2.f32 byte0, %6, %5;\n"
"cvt.rn.satfinite.e2m1x2.f32 byte1, %8, %7;\n"
"cvt.rn.satfinite.e2m1x2.f32 byte2, %10, %9;\n"
"cvt.rn.satfinite.e2m1x2.f32 byte3, %12, %11;\n"
"mov.b32 %0, {byte0, byte1, byte2, byte3};\n"
"cvt.rn.f16x2.e2m1x2 %1, byte0;\n"
"cvt.rn.f16x2.e2m1x2 %2, byte1;\n"
"cvt.rn.f16x2.e2m1x2 %3, byte2;\n"
"cvt.rn.f16x2.e2m1x2 %4, byte3;\n"
"}"
: "=r"(out), "=r"(out_dequant_1), "=r"(out_dequant_2), "=r"(out_dequant_3), "=r"(out_dequant_4) : "f"(x_scaled[0]), "f"(x_scaled[1]), "f"(x_scaled[2]), "f"(x_scaled[3]),
"f"(x_scaled[4]), "f"(x_scaled[5]), "f"(x_scaled[6]), "f"(x_scaled[7]));
unsigned short out_dequant_1_hi = (out_dequant_1 >> 16) & 0xFFFF;
unsigned short out_dequant_1_lo = out_dequant_1 & 0xFFFF;
unsigned short out_dequant_2_hi = (out_dequant_2 >> 16) & 0xFFFF;
unsigned short out_dequant_2_lo = out_dequant_2 & 0xFFFF;
unsigned short out_dequant_3_hi = (out_dequant_3 >> 16) & 0xFFFF;
unsigned short out_dequant_3_lo = out_dequant_3 & 0xFFFF;
unsigned short out_dequant_4_hi = (out_dequant_4 >> 16) & 0xFFFF;
unsigned short out_dequant_4_lo = out_dequant_4 & 0xFFFF;
float val0 = __half2float(__ushort_as_half(out_dequant_1_lo)) * sf * amax / (E2M1_MAX_VALUE * E4M3_MAX_FOUROVERSIX);
float val1 = __half2float(__ushort_as_half(out_dequant_1_hi)) * sf * amax / (E2M1_MAX_VALUE * E4M3_MAX_FOUROVERSIX);
float val2 = __half2float(__ushort_as_half(out_dequant_2_lo)) * sf * amax / (E2M1_MAX_VALUE * E4M3_MAX_FOUROVERSIX);
float val3 = __half2float(__ushort_as_half(out_dequant_2_hi)) * sf * amax / (E2M1_MAX_VALUE * E4M3_MAX_FOUROVERSIX);
float val4 = __half2float(__ushort_as_half(out_dequant_3_lo)) * sf * amax / (E2M1_MAX_VALUE * E4M3_MAX_FOUROVERSIX);
float val5 = __half2float(__ushort_as_half(out_dequant_3_hi)) * sf * amax / (E2M1_MAX_VALUE * E4M3_MAX_FOUROVERSIX);
float val6 = __half2float(__ushort_as_half(out_dequant_4_lo)) * sf * amax / (E2M1_MAX_VALUE * E4M3_MAX_FOUROVERSIX);
float val7 = __half2float(__ushort_as_half(out_dequant_4_hi)) * sf * amax / (E2M1_MAX_VALUE * E4M3_MAX_FOUROVERSIX);
if constexpr (kAdaptiveBlockScalingRuleType == AdaptiveBlockScalingRuleType::MAE_4o6)
{
*err += std::abs(val0 - x[0]);
*err += std::abs(val1 - x[1]);
*err += std::abs(val2 - x[2]);
*err += std::abs(val3 - x[3]);
*err += std::abs(val4 - x[4]);
*err += std::abs(val5 - x[5]);
*err += std::abs(val6 - x[6]);
*err += std::abs(val7 - x[7]);
}
else if constexpr (kAdaptiveBlockScalingRuleType == AdaptiveBlockScalingRuleType::MSE_4o6)
{
*err += (val0 - x[0]) * (val0 - x[0]);
*err += (val1 - x[1]) * (val1 - x[1]);
*err += (val2 - x[2]) * (val2 - x[2]);
*err += (val3 - x[3]) * (val3 - x[3]);
*err += (val4 - x[4]) * (val4 - x[4]);
*err += (val5 - x[5]) * (val5 - x[5]);
*err += (val6 - x[6]) * (val6 - x[6]);
*err += (val7 - x[7]) * (val7 - x[7]);
}
else if constexpr (kAdaptiveBlockScalingRuleType == AdaptiveBlockScalingRuleType::ABS_MAX_4o6)
{
float val0_err = std::abs(val0 - x[0]);
if (val0_err > *err)
*err = val0_err;
float val1_err = std::abs(val1 - x[1]);
if (val1_err > *err)
*err = val1_err;
float val2_err = std::abs(val2 - x[2]);
if (val2_err > *err)
*err = val2_err;
float val3_err = std::abs(val3 - x[3]);
if (val3_err > *err)
*err = val3_err;
float val4_err = std::abs(val4 - x[4]);
if (val4_err > *err)
*err = val4_err;
float val5_err = std::abs(val5 - x[5]);
if (val5_err > *err)
*err = val5_err;
float val6_err = std::abs(val6 - x[6]);
if (val6_err > *err)
*err = val6_err;
float val7_err = std::abs(val7 - x[7]);
if (val7_err > *err)
*err = val7_err;
}
else
{
printf("in Fp4ArrayQuant::convert, kAdaptiveBlockScalingRuleType = %d, not supported\n", kAdaptiveBlockScalingRuleType);
assert(false);
}
return reinterpret_cast<OutputType const &>(out);
}
else
{
asm volatile(
"{\n"
".reg .b8 byte0, byte1, byte2, byte3;\n"
"cvt.rn.satfinite.e2m1x2.f32 byte0, %2, %1;\n"
"cvt.rn.satfinite.e2m1x2.f32 byte1, %4, %3;\n"
"cvt.rn.satfinite.e2m1x2.f32 byte2, %6, %5;\n"
"cvt.rn.satfinite.e2m1x2.f32 byte3, %8, %7;\n"
"mov.b32 %0, {byte0, byte1, byte2, byte3};\n"
"}"
: "=r"(out) : "f"(x_scaled[0]), "f"(x_scaled[1]), "f"(x_scaled[2]), "f"(x_scaled[3]),
"f"(x_scaled[4]), "f"(x_scaled[5]), "f"(x_scaled[6]), "f"(x_scaled[7]));
return reinterpret_cast<OutputType const &>(out);
}
}
else
{
if constexpr (Is_4o6)
{
unsigned out_dequant_1;
unsigned out_dequant_2;
unsigned out_dequant_3;
unsigned out_dequant_4;
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ == 1000 || __CUDA_ARCH__ == 1030)
asm volatile(
"{\n"
".reg .b16 tmp0, tmp1;\n"
".reg .b8 byte0, byte1;\n"
"cvt.rs.satfinite.e2m1x4.f32 tmp0, {%8, %7, %6, %5}, %13;\n"
"mov.b16 {byte1, byte0}, tmp0;\n"
"cvt.rn.f16x2.e2m1x2 %1, byte0;\n"
"cvt.rn.f16x2.e2m1x2 %2, byte1;\n"
"cvt.rs.satfinite.e2m1x4.f32 tmp1, {%12, %11, %10, %9}, %14;\n"
"mov.b16 {byte1, byte0}, tmp1;\n"
"cvt.rn.f16x2.e2m1x2 %3, byte0;\n"
"cvt.rn.f16x2.e2m1x2 %4, byte1;\n"
"mov.b32 %0, {tmp0, tmp1};\n"
"}"
: "=r"(out), "=r"(out_dequant_1), "=r"(out_dequant_2), "=r"(out_dequant_3), "=r"(out_dequant_4) : "f"(x_scaled[0]), "f"(x_scaled[1]), "f"(x_scaled[2]), "f"(x_scaled[3]), "f"(x_scaled[4]), "f"(x_scaled[5]), "f"(x_scaled[6]), "f"(x_scaled[7]), "r"(rbits), "r"(rbits));
#endif
unsigned short out_dequant_1_hi = (out_dequant_1 >> 16) & 0xFFFF;
unsigned short out_dequant_1_lo = out_dequant_1 & 0xFFFF;
unsigned short out_dequant_2_hi = (out_dequant_2 >> 16) & 0xFFFF;
unsigned short out_dequant_2_lo = out_dequant_2 & 0xFFFF;
unsigned short out_dequant_3_hi = (out_dequant_3 >> 16) & 0xFFFF;
unsigned short out_dequant_3_lo = out_dequant_3 & 0xFFFF;
unsigned short out_dequant_4_hi = (out_dequant_4 >> 16) & 0xFFFF;
unsigned short out_dequant_4_lo = out_dequant_4 & 0xFFFF;
float val0 = __half2float(__ushort_as_half(out_dequant_1_lo)) * sf * amax / (E2M1_MAX_VALUE * E4M3_MAX_FOUROVERSIX);
float val1 = __half2float(__ushort_as_half(out_dequant_1_hi)) * sf * amax / (E2M1_MAX_VALUE * E4M3_MAX_FOUROVERSIX);
float val2 = __half2float(__ushort_as_half(out_dequant_2_lo)) * sf * amax / (E2M1_MAX_VALUE * E4M3_MAX_FOUROVERSIX);
float val3 = __half2float(__ushort_as_half(out_dequant_2_hi)) * sf * amax / (E2M1_MAX_VALUE * E4M3_MAX_FOUROVERSIX);
float val4 = __half2float(__ushort_as_half(out_dequant_3_lo)) * sf * amax / (E2M1_MAX_VALUE * E4M3_MAX_FOUROVERSIX);
float val5 = __half2float(__ushort_as_half(out_dequant_3_hi)) * sf * amax / (E2M1_MAX_VALUE * E4M3_MAX_FOUROVERSIX);
float val6 = __half2float(__ushort_as_half(out_dequant_4_lo)) * sf * amax / (E2M1_MAX_VALUE * E4M3_MAX_FOUROVERSIX);
float val7 = __half2float(__ushort_as_half(out_dequant_4_hi)) * sf * amax / (E2M1_MAX_VALUE * E4M3_MAX_FOUROVERSIX);
if constexpr (kAdaptiveBlockScalingRuleType == AdaptiveBlockScalingRuleType::MAE_4o6)
{
*err += std::abs(val0 - x[0]);
*err += std::abs(val1 - x[1]);
*err += std::abs(val2 - x[2]);
*err += std::abs(val3 - x[3]);
*err += std::abs(val4 - x[4]);
*err += std::abs(val5 - x[5]);
*err += std::abs(val6 - x[6]);
*err += std::abs(val7 - x[7]);
}
else if constexpr (kAdaptiveBlockScalingRuleType == AdaptiveBlockScalingRuleType::MSE_4o6)
{
*err += (val0 - x[0]) * (val0 - x[0]);
*err += (val1 - x[1]) * (val1 - x[1]);
*err += (val2 - x[2]) * (val2 - x[2]);
*err += (val3 - x[3]) * (val3 - x[3]);
*err += (val4 - x[4]) * (val4 - x[4]);
*err += (val5 - x[5]) * (val5 - x[5]);
*err += (val6 - x[6]) * (val6 - x[6]);
*err += (val7 - x[7]) * (val7 - x[7]);
}
else if constexpr (kAdaptiveBlockScalingRuleType == AdaptiveBlockScalingRuleType::ABS_MAX_4o6)
{
float val0_err = std::abs(val0 - x[0]);
if (val0_err > *err)
*err = val0_err;
float val1_err = std::abs(val1 - x[1]);
if (val1_err > *err)
*err = val1_err;
float val2_err = std::abs(val2 - x[2]);
if (val2_err > *err)
*err = val2_err;
float val3_err = std::abs(val3 - x[3]);
if (val3_err > *err)
*err = val3_err;
float val4_err = std::abs(val4 - x[4]);
if (val4_err > *err)
*err = val4_err;
float val5_err = std::abs(val5 - x[5]);
if (val5_err > *err)
*err = val5_err;
float val6_err = std::abs(val6 - x[6]);
if (val6_err > *err)
*err = val6_err;
float val7_err = std::abs(val7 - x[7]);
if (val7_err > *err)
*err = val7_err;
}
else
{
printf("in Fp4ArrayQuant::convert, kAdaptiveBlockScalingRuleType = %d, not supported\n", kAdaptiveBlockScalingRuleType);
assert(false);
}
return reinterpret_cast<OutputType const &>(out);
}
else
{
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ == 1000 || __CUDA_ARCH__ == 1030)
asm volatile(
"{\n"
".reg .b16 tmp0, tmp1;\n"
"cvt.rs.satfinite.e2m1x4.f32 tmp0, {%4, %3, %2, %1}, %9;\n"
"cvt.rs.satfinite.e2m1x4.f32 tmp1, {%8, %7, %6, %5}, %10;\n"
"mov.b32 %0, {tmp0, tmp1};\n"
"}"
: "=r"(out) : "f"(x_scaled[0]), "f"(x_scaled[1]), "f"(x_scaled[2]), "f"(x_scaled[3]),
"f"(x_scaled[4]), "f"(x_scaled[5]), "f"(x_scaled[6]), "f"(x_scaled[7]),
"r"(rbits), "r"(rbits));
#endif
return reinterpret_cast<OutputType const &>(out);
}
}
}
};
template <bool Is_nvfp4, bool Is_2d, bool Is_4o6, bool Is_rtn, AdaptiveBlockScalingRuleType kRule, typename Engine, typename Layout, typename OutputType>
__forceinline__ __device__ float fp4_conversion(Tensor<Engine, Layout> const &tensor, const float amax, float *sf_, OutputType *res, const uint32_t rbits)
{
constexpr int numel = decltype(size(tensor))::value;
static_assert((numel == 16 && Is_nvfp4) || numel == 32);
static_assert(std::is_same_v<OutputType, cutlass::Array<cutlass::float_e2m1_t, 8>>);
constexpr int loop_size = 8;
constexpr int num_loops = numel / loop_size;
using InputType = cutlass::Array<float, loop_size>;
Fp4ArrayQuant<Is_4o6, Is_rtn, kRule> fp4_array_quant;
if constexpr (Is_4o6)
{
float err[2] = {0.0f, 0.0f};
OutputType res_4[num_loops];
OutputType res_6[num_loops];
float final_err[2] = {0.0f, 0.0f};
#pragma unroll
for (int i = 0; i < num_loops; ++i)
{
InputType x;
#pragma unroll
for (int j = 0; j < loop_size; ++j)
{
x[j] = static_cast<float>(tensor(i * loop_size + j));
}
res_4[i] = fp4_array_quant.convert(x, amax, sf_[0], rbits, &err[0]);
res_6[i] = fp4_array_quant.convert(x, amax, sf_[1], rbits, &err[1]);
}
if (Is_2d){
// For 2D tensors we want to pick the same format for the entire tensor, to keep it simple for downstream processing.
// So we pick the format with smaller total error across the entire tensor.
// If the method is MAE or MSE, we need to sum the error across the entire tensor. If the method is ABS_MAX, we need to take the max error across the entire tensor.
using RedOp = std::conditional_t<kRule == AdaptiveBlockScalingRuleType::ABS_MAX_4o6, MaxOp<float>, SumOp<float>>;
RedOp op;
final_err[0] = Allreduce<numel>::run(err[0], op);
final_err[1] = Allreduce<numel>::run(err[1], op);
} else {
final_err[0] = err[0];
final_err[1] = err[1];
}
// pick_first = true means choose 4, false means choose 6
bool const pick_first = final_err[0] < final_err[1];
#pragma unroll
for (int i = 0; i < num_loops; ++i)
{
res[i] = pick_first ? res_4[i] : res_6[i];
}
return sf_[!pick_first];
}
else
{
#pragma unroll
for (int i = 0; i < num_loops; ++i)
{
InputType x;
#pragma unroll
for (int j = 0; j < loop_size; ++j)
{
x[j] = static_cast<float>(tensor(i * loop_size + j));
}
res[i] = fp4_array_quant.convert(x, amax, sf_[0], rbits, nullptr);
}
return sf_[0];
}
}
__device__ __forceinline__ float clamp(float value, float min_value, float max_value)
{
return max(min(value, max_value), min_value);
}
////////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace fouroversix
fouroversix/src/fouroversix/csrc/quantize/fp4_quant.cu
+245
View File
@@ -0,0 +1,245 @@
#include <pybind11/pybind11.h>
#include <torch/extension.h>
#include <torch/python.h>
#include <torch/nn/functional.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAStream.h>
#include <ATen/cuda/CUDAGeneratorImpl.h> // For at::Generator and at::PhiloxCudaState
#include <cutlass/numeric_types.h>
#include "hardware_info.h"
#include "fp4_quant.h"
#include "static_switch.h"
#define CHECK_DEVICE(x) TORCH_CHECK(x.is_cuda(), #x " must be on CUDA")
#define CHECK_SHAPE(x, ...) TORCH_CHECK(x.sizes() == torch::IntArrayRef({__VA_ARGS__}), #x " must have shape (" #__VA_ARGS__ ")")
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
namespace fouroversix
{
void set_params_fp4_quant(
FP4_quant_params &params,
/*-------------- tensors ---------------*/
const at::Tensor x,
at::Tensor x_rht,
at::Tensor x_e2m1,
at::Tensor x_sf,
at::Tensor x_sft,
at::Tensor amax,
const int M,
const int N,
const int M_rounded,
const int N_rounded,
const int M_sf,
const int N_sf,
const bool is_nvfp4,
const bool is_rtn,
const bool is_rht,
const bool is_2d,
// const bool is_4o6,
const bool is_transpose,
const int selection_rule,
const int rbits)
{
// Reset the parameters
params = {};
params.is_bf16 = x.dtype() == torch::kBFloat16;
/**************** Pointers & strides ****************/
params.x_ptr = x.data_ptr();
params.x_rht_ptr = x_rht.data_ptr();
params.x_e2m1_ptr = x_e2m1.data_ptr();
params.x_sf_ptr = x_sf.data_ptr();
params.x_sft_ptr = x_sft.data_ptr();
params.amax_ptr = amax.data_ptr();
// Element-based strides (not bytes)
params.x_row_stride = x.stride(0);
params.x_col_stride = x.stride(1);
params.x_rht_row_stride = x_rht.stride(0);
params.x_rht_col_stride = x_rht.stride(1);
params.x_e2m1_row_stride = x_e2m1.stride(0);
params.x_e2m1_col_stride = x_e2m1.stride(1);
params.x_sf_row_stride = x_sf.stride(0);
params.x_sf_col_stride = x_sf.stride(1);
params.x_sft_row_stride = x_sft.stride(0);
params.x_sft_col_stride = x_sft.stride(1);
// Set the dimensions
params.M = M;
params.N = N;
params.M_rounded = M_rounded;
params.N_rounded = N_rounded;
params.M_sf = M_sf;
params.N_sf = N_sf;
// Set FP4-specific parameters
params.is_nvfp4 = is_nvfp4;
params.is_rtn = is_rtn;
params.is_rht = is_rht;
params.is_2d = is_2d;
// params.is_4o6 = is_4o6;
params.is_transpose = is_transpose;
params.selection_rule = selection_rule;
params.rbits = rbits;
}
void run_fp4_quant(FP4_quant_params &params, cudaStream_t stream)
{
FP16_SWITCH(!params.is_bf16, [&]
{ BOOL_SWITCH(params.is_nvfp4, Is_nvfp4, [&]
{ BOOL_SWITCH(params.is_rht, Is_rht, [&]
{ BOOL_SWITCH(params.is_transpose, Is_transpose, [&]
{ run_fp4_quant_<fp16_type, Is_nvfp4, Is_rht, Is_transpose>(params, stream); }); }); }); });
}
std::tuple<at::Tensor, at::Tensor, at::Tensor> quantize_to_fp4(
const at::Tensor &x,
const bool is_nvfp4,
const bool is_rtn,
const bool is_rht,
// const bool is_4o6,
const bool is_2d,
const bool is_transpose,
const int64_t selection_rule,
const int64_t rbits)
{
/*******
* selection_rule:
* 0: static_6
* 1: static_4
* 2: 4o6_l1_norm
* 3: 4o6_mse
* 4: 4o6_abs_max
*/
TORCH_CHECK(selection_rule >= 0 && selection_rule <= 4, "Invalid selection_rule: " + std::to_string(selection_rule));
// const int is_4o6 = selection_rule == 2 || selection_rule == 3;
/**********************
* 1. Sanity checks *
*********************/
at::cuda::CUDAGuard device_guard{x.device()};
// Hardware capability
{
auto [cc_major, _] = get_compute_capability(get_current_device());
TORCH_CHECK(cc_major >= 10, "FP4Quant only supports Blackwell GPUs or newer.");
}
// Dtype / device checks
auto dtype = x.dtype();
TORCH_CHECK(dtype == torch::kFloat16 || dtype == torch::kBFloat16,
"FP4Quant only supports fp16 and bf16 data types");
// TORCH_CHECK(km.dtype() == dtype, "q and km must have the same dtype");
CHECK_DEVICE(x);
// Layout / contiguity checks
TORCH_CHECK(x.stride(-1) == 1, "x must be contiguous on the last dim");
/**********************
* 2. Dimension logic *
*********************/
const int M = is_transpose ? x.size(1) : x.size(0);
const int N = is_transpose ? x.size(0) : x.size(1);
const int n_round = is_nvfp4 ? 64 : 128;
TORCH_CHECK(N % n_round == 0, "N must be multiple of " + std::to_string(n_round));
/**********************
* 3. Derived sizes *
*********************/
auto round_up = [](int x, int m)
{ return (x + m - 1) / m * m; };
const int M_rounded = round_up(M, 128);
const int N_rounded = round_up(N, n_round);
// const int max_seqlen_km = ceil_div(max_seqlen_k, moba_chunk_size);
// const int head_size_rounded = round_up(head_size, head_size <= 128 ? 32 : 64);
// const int seqlen_q_rounded = round_up(max_seqlen_q, 128);
// const int seqlen_km_rounded = round_up(max_seqlen_km, 128);
// const int moba_topk_rounded = round_up(moba_topk, 16);
/**********************
* 4. Intermediate buffers *
*********************/
at::Tensor x_rht;
if (is_rht)
{
x_rht = torch::zeros({M_rounded, N_rounded}, x.options());
}
else
{
x_rht = torch::zeros({0, 0}, x.options());
}
/**********************
* 5. Output buffers *
*********************/
at::Tensor x_e2m1 = torch::zeros({M_rounded, int(N_rounded / 2)}, x.options().dtype(torch::kUInt8));
at::Tensor x_sf, x_sft;
int M_sf, N_sf;
if (is_nvfp4)
{
M_sf = int(M_rounded / 128 * 32) * int(N_rounded / 64);
N_sf = 16;
// N_sf = int(N_rounded / 16 * 4);
x_sf = torch::zeros({M_sf, N_sf}, x.options().dtype(torch::kFloat8_e4m3fn));
x_sft = torch::zeros({M_rounded, int(N_rounded / 16)}, x.options().dtype(torch::kFloat32));
}
else
{
M_sf = int(M_rounded / 128 * 32) * int(N_rounded / 128);
N_sf = 16;
x_sf = torch::zeros({M_sf, N_sf}, x.options().dtype(torch::kUInt8));
x_sft = torch::zeros({M_rounded, int(N_rounded / 32)}, x.options().dtype(torch::kFloat32));
}
at::Tensor amax = torch::zeros({1}, x.options().dtype(torch::kFloat32));
/**********************
* 5. Param struct *
*********************/
FP4_quant_params params;
// const at::Tensor x,
// at::Tensor x_rht,
// at::Tensor x_e2m1,
// at::Tensor x_sf,
// at::Tensor amax,
// const int M,
// const int N,
// const bool is_nvfp4,
// const bool is_rtn,
// const bool is_4o6,
// const bool is_2d,
// const bool is_transpose
set_params_fp4_quant(
params,
/*-------------- tensors ---------------*/
x, x_rht, x_e2m1, x_sf, x_sft, amax, M, N, M_rounded, N_rounded, M_sf, N_sf,
is_nvfp4, is_rtn, is_rht, is_2d, is_transpose, selection_rule, rbits);
/**********************
* 6. Kernel launch *
*********************/
if (M > 0)
{
run_fp4_quant(params, at::cuda::getCurrentCUDAStream().stream());
}
else
{
amax.fill_(0);
}
return std::make_tuple(x_e2m1, x_sf.flatten(), amax);
}
TORCH_LIBRARY_IMPL(fouroversix, CUDA, m)
{
m.impl("quantize_to_fp4", &quantize_to_fp4);
}
}
fouroversix/src/fouroversix/csrc/quantize/fp4_quant_bf16_mxfp4_rht_sm100.cu
+11
View File
@@ -0,0 +1,11 @@
// Splitting the different transpose modes to different files to speed up compilation.
// This file is auto-generated. See "generate_kernels.py"
#include "fp4_quant_launch_template.h"
namespace fouroversix {
template<>
void run_fp4_quant_<cutlass::bfloat16_t, false, true, false>(FP4_quant_params &params, cudaStream_t stream) {
run_mxfp4_quant_rht<cutlass::bfloat16_t, false>(params, stream);
}
} // namespace fouroversix
fouroversix/src/fouroversix/csrc/quantize/fp4_quant_bf16_mxfp4_rht_trans_sm100.cu
+11
View File
@@ -0,0 +1,11 @@
// Splitting the different transpose modes to different files to speed up compilation.
// This file is auto-generated. See "generate_kernels.py"
#include "fp4_quant_launch_template.h"
namespace fouroversix {
template<>
void run_fp4_quant_<cutlass::bfloat16_t, false, true, true>(FP4_quant_params &params, cudaStream_t stream) {
run_mxfp4_quant_rht<cutlass::bfloat16_t, true>(params, stream);
}
} // namespace fouroversix
fouroversix/src/fouroversix/csrc/quantize/fp4_quant_bf16_mxfp4_sm100.cu
+11
View File
@@ -0,0 +1,11 @@
// Splitting the different transpose modes to different files to speed up compilation.
// This file is auto-generated. See "generate_kernels.py"
#include "fp4_quant_launch_template.h"
namespace fouroversix {
template<>
void run_fp4_quant_<cutlass::bfloat16_t, false, false, false>(FP4_quant_params &params, cudaStream_t stream) {
run_mxfp4_quant<cutlass::bfloat16_t, false>(params, stream);
}
} // namespace fouroversix
fouroversix/src/fouroversix/csrc/quantize/fp4_quant_bf16_mxfp4_trans_sm100.cu
+11
View File
@@ -0,0 +1,11 @@
// Splitting the different transpose modes to different files to speed up compilation.
// This file is auto-generated. See "generate_kernels.py"
#include "fp4_quant_launch_template.h"
namespace fouroversix {
template<>
void run_fp4_quant_<cutlass::bfloat16_t, false, false, true>(FP4_quant_params &params, cudaStream_t stream) {
run_mxfp4_quant<cutlass::bfloat16_t, true>(params, stream);
}
} // namespace fouroversix
fouroversix/src/fouroversix/csrc/quantize/fp4_quant_bf16_nvfp4_rht_sm100.cu
+11
View File
@@ -0,0 +1,11 @@
// Splitting the different transpose modes to different files to speed up compilation.
// This file is auto-generated. See "generate_kernels.py"
#include "fp4_quant_launch_template.h"
namespace fouroversix {
template<>
void run_fp4_quant_<cutlass::bfloat16_t, true, true, false>(FP4_quant_params &params, cudaStream_t stream) {
run_nvfp4_quant_rht<cutlass::bfloat16_t, false>(params, stream);
}
} // namespace fouroversix
fouroversix/src/fouroversix/csrc/quantize/fp4_quant_bf16_nvfp4_rht_trans_sm100.cu
+11
View File
@@ -0,0 +1,11 @@
// Splitting the different transpose modes to different files to speed up compilation.
// This file is auto-generated. See "generate_kernels.py"
#include "fp4_quant_launch_template.h"
namespace fouroversix {
template<>
void run_fp4_quant_<cutlass::bfloat16_t, true, true, true>(FP4_quant_params &params, cudaStream_t stream) {
run_nvfp4_quant_rht<cutlass::bfloat16_t, true>(params, stream);
}
} // namespace fouroversix

Some files were not shown because too many files have changed in this diff Show More