diff --git a/.github/workflows/fastvideo-kernel-publish.yml b/.github/workflows/fastvideo-kernel-publish.yml new file mode 100644 index 000000000..4a1e4fa44 --- /dev/null +++ b/.github/workflows/fastvideo-kernel-publish.yml @@ -0,0 +1,236 @@ +name: Publish FastVideo Kernel to PyPI on Version Change + +on: + push: + branches: + - main + paths: + - "csrc/fastvideo_kernel/pyproject.toml" + workflow_dispatch: + +jobs: + check-version-change: + runs-on: ubuntu-latest + outputs: + version-changed: ${{ steps.check-version.outputs.changed }} + new-version: ${{ steps.check-version.outputs.new-version }} + steps: + - name: Checkout code + uses: actions/checkout@v4 + with: + fetch-depth: 2 + + - name: Check if version changed + id: check-version + run: | + cd csrc/fastvideo_kernel + # Get current commit's version from pyproject.toml + NEW_VERSION=$(grep -oP 'version\s*=\s*"\K[^"]+' pyproject.toml) + echo "New version: $NEW_VERSION" + + # Get previous version from git history + OLD_VERSION=$(git show HEAD~1:./pyproject.toml | grep -oP 'version\s*=\s*"\K[^"]+' || echo "0.0.0") + echo "Old version: $OLD_VERSION" + + if [ "$NEW_VERSION" != "$OLD_VERSION" ]; then + echo "Version changed from $OLD_VERSION to $NEW_VERSION" + echo "changed=true" >> $GITHUB_OUTPUT + echo "new-version=$NEW_VERSION" >> $GITHUB_OUTPUT + else + echo "Version did not change" + echo "changed=false" >> $GITHUB_OUTPUT + fi + + build_wheels: + name: Build Wheel + needs: check-version-change + if: ${{ needs.check-version-change.outputs.version-changed == 'true' || github.event_name == 'workflow_dispatch' }} + runs-on: ${{ matrix.os }} + + strategy: + fail-fast: false + matrix: + os: [ubuntu-22.04] + python-version: ['3.10', '3.11', '3.12', '3.13'] + torch-cuda: + - torch-version: '2.5.1' + cuda-version: '12.4.1' + torch-cuda-short: 'cu124' + - torch-version: '2.6.0' + cuda-version: '12.6.3' + torch-cuda-short: 'cu126' + - torch-version: '2.7.1' + cuda-version: '12.8.0' + torch-cuda-short: 'cu128' + + steps: + - name: Free up disk space + run: | + echo "Initial disk space:" + df -h + + # Remove large directories + sudo rm -rf /usr/share/dotnet + sudo rm -rf /usr/local/lib/android + sudo rm -rf /opt/ghc + sudo rm -rf /usr/local/share/boost + sudo rm -rf /usr/share/swift + sudo rm -rf /usr/local/lib/node_modules + sudo rm -rf /usr/local/share/powershell + sudo rm -rf /usr/share/rust + sudo rm -rf /usr/local/.ghcup + + # Remove cached files + sudo rm -rf /var/lib/apt/lists/* + sudo rm -rf /var/cache/apt/archives/* + + echo "Disk space after cleanup:" + df -h + + - name: Checkout + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + + - name: Install CUDA ${{ matrix.torch-cuda.cuda-version }} + uses: Jimver/cuda-toolkit@v0.2.21 + id: cuda-toolkit + with: + cuda: ${{ matrix.torch-cuda.cuda-version }} + linux-local-args: '["--toolkit"]' + method: 'network' + + - name: Install dependencies (GCC, Clang, CUDA Paths, Git) + run: | + sudo apt update + sudo apt install -y git patchelf gcc-11 g++-11 clang-11 + sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11 + + # Allow Git to Access Safe Directory + git config --global --add safe.directory /__w/FastVideo/FastVideo + + # Set CUDA environment variables + export CUDA_HOME=/usr/local/cuda-${{ matrix.torch-cuda.cuda-version }} + export PATH=${CUDA_HOME}/bin:${PATH} + export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH + + # Verify installation + gcc --version + g++ --version + clang-11 --version + nvcc --version + + - name: Install PyTorch ${{ matrix.torch-cuda.torch-version }}+cu${{ matrix.torch-cuda.cuda-version }} + run: | + pip install --upgrade pip + pip install typing-extensions==4.12.2 + pip install --no-cache-dir torch==${{ matrix.torch-cuda.torch-version }} --index-url https://download.pytorch.org/whl/${{matrix.torch-cuda.torch-cuda-short}} + 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: Build wheel + run: | + export PYTHONPATH=$GITHUB_WORKSPACE:$PYTHONPATH + + pip install setuptools ninja packaging wheel triton + + cd csrc/fastvideo_kernel + git submodule update --init --recursive # Ensure ThunderKittens submodule is initialized + python setup.py bdist_wheel --dist-dir=dist + + - name: Rename wheel file + run: | + cd csrc/fastvideo_kernel + + CUDA_SHORT_VERSION=$(echo ${{ matrix.torch-cuda.cuda-version }} | cut -d. -f1,2 | sed 's/\.//g') + TORCH_SHORT_VERSION=$(echo ${{ matrix.torch-cuda.torch-version }} | cut -d. -f1,2) + # Get the correct version format + tmpname=cu${CUDA_SHORT_VERSION}torch${TORCH_SHORT_VERSION} + wheel_name=$(ls dist/*whl | xargs -n 1 basename | sed "s/-/+$tmpname-/2") + # Rename with version information + ls dist/*whl |xargs -I {} mv {} dist/${wheel_name} + echo "wheel_name=${wheel_name}" >> $GITHUB_ENV + + - name: Upload wheel artifact + uses: actions/upload-artifact@v4 + with: + name: ${{ env.wheel_name }}-py${{ matrix.python-version }} + path: csrc/fastvideo_kernel/dist/*.whl + retention-days: 90 + + publish_package: + name: Publish package + needs: [build_wheels, check-version-change] + if: ${{ needs.check-version-change.outputs.version-changed == 'true' || github.event_name == 'workflow_dispatch' }} + runs-on: ubuntu-22.04 + permissions: + id-token: write # Needed for OIDC Trusted Publishing + + steps: + - uses: actions/checkout@v4 + + - uses: actions/setup-python@v5 + with: + python-version: '3.10' + + - name: Install CUDA 12.4.1 + uses: Jimver/cuda-toolkit@v0.2.21 + id: cuda-toolkit + with: + cuda: 12.4.1 + linux-local-args: '["--toolkit"]' + method: 'network' + sub-packages: '["nvcc"]' + + - name: Install dependencies (GCC, Clang, CUDA Paths, Git) + run: | + sudo apt update + sudo apt install -y git patchelf gcc-11 g++-11 clang-11 + sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11 + + # Allow Git to Access Safe Directory + git config --global --add safe.directory /__w/FastVideo/FastVideo + + # Set CUDA environment variables + export CUDA_HOME=/usr/local/cuda-12.4.1 + export PATH=${CUDA_HOME}/bin:${PATH} + export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH + + # Verify installation + gcc --version + g++ --version + clang-11 --version + nvcc --version + + - name: Install PyTorch 2.5.1+cu12.4.1 + run: | + pip install --upgrade pip + pip install typing-extensions==4.12.2 + export TORCH_CUDA_VERSION=124 + pip install --no-cache-dir torch==2.5.1 --index-url https://download.pytorch.org/whl/cu${TORCH_CUDA_VERSION} + 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: Build source distribution + run: | + export PYTHONPATH=$GITHUB_WORKSPACE:$PYTHONPATH + + pip install setuptools ninja packaging wheel triton + + cd csrc/fastvideo_kernel + git submodule update --init --recursive + python setup.py sdist --dist-dir=dist + + - name: Publish release distributions to PyPI + uses: pypa/gh-action-pypi-publish@release/v1 + with: + packages-dir: csrc/fastvideo_kernel/dist/ diff --git a/csrc/fastvideo_kernel/.gitignore b/csrc/fastvideo_kernel/.gitignore new file mode 100644 index 000000000..317ddafc4 --- /dev/null +++ b/csrc/fastvideo_kernel/.gitignore @@ -0,0 +1,7 @@ +build/ +dist/ +*.egg-info/ +__pycache__/ +*.so +*.pyc +.ipynb_checkpoints/ diff --git a/csrc/fastvideo_kernel/LICENSE b/csrc/fastvideo_kernel/LICENSE new file mode 100644 index 000000000..a7b8b5e6a --- /dev/null +++ b/csrc/fastvideo_kernel/LICENSE @@ -0,0 +1,187 @@ + 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. \ No newline at end of file diff --git a/csrc/fastvideo_kernel/MANIFEST.in b/csrc/fastvideo_kernel/MANIFEST.in new file mode 100644 index 000000000..6d2e724a1 --- /dev/null +++ b/csrc/fastvideo_kernel/MANIFEST.in @@ -0,0 +1,6 @@ +include LICENSE +include README.md +include pyproject.toml +recursive-include src/fastvideo_kernel *.cu *.cuh *.cpp *.h +recursive-include csrc *.cu *.cuh *.cpp *.h +recursive-include tk *.cu *.cuh *.cpp *.h diff --git a/csrc/fastvideo_kernel/README.md b/csrc/fastvideo_kernel/README.md new file mode 100644 index 000000000..72412db3f --- /dev/null +++ b/csrc/fastvideo_kernel/README.md @@ -0,0 +1,31 @@ +# FastVideo Kernel + +CUDA kernels for FastVideo video generation. + +## Installation + +```bash +git submodule update --init --recursive +cd csrc/fastvideo_kernel +pip install . +``` + +## Usage + +```python +from fastvideo_kernel import sliding_tile_attention, video_sparse_attn, moba_attn_varlen + +# Example: Sliding Tile Attention +out = sliding_tile_attention(q, k, v, window_sizes, text_len) + +# Example: Video Sparse Attention (with Triton fallback) +out = video_sparse_attn(q, k, v, block_sizes, topk=5) + +# Example: VMoBA +out = moba_attn_varlen(q, k, v, cu_seqlens_q, cu_seqlens_k, ...) +``` + +## Requirements + +- H100 GPU (sm_90a) for CUDA kernels +- Triton for non-H100 fallback diff --git a/csrc/fastvideo_kernel/csrc/block_sparse_h100.cu b/csrc/fastvideo_kernel/csrc/block_sparse_h100.cu new file mode 100644 index 000000000..58031322f --- /dev/null +++ b/csrc/fastvideo_kernel/csrc/block_sparse_h100.cu @@ -0,0 +1,1236 @@ +// # Define TORCH_COMPILE macro + +#include "kittens.cuh" +#include +#include +#include + + +using namespace kittens; +namespace cg = cooperative_groups; +constexpr int BLOCK_M = 64; +constexpr int BLOCK_N = 64; +template struct fwd_attend_ker_tile_dims {}; +template<> struct fwd_attend_ker_tile_dims<64> { + constexpr static int tile_width = (64); + constexpr static int qo_height = (4*16); + constexpr static int kv_height = (4*16); +}; +template<> struct fwd_attend_ker_tile_dims<128> { + constexpr static int tile_width = (128); + constexpr static int qo_height = (4*16); + constexpr static int kv_height = (4*16); +}; +template struct fwd_globals { + using q_tile = st_bf::qo_height, fwd_attend_ker_tile_dims::tile_width>; + using k_tile = st_bf::kv_height, fwd_attend_ker_tile_dims::tile_width>; + using v_tile = st_bf::kv_height, fwd_attend_ker_tile_dims::tile_width>; + using l_col_vec = col_vec::qo_height, fwd_attend_ker_tile_dims::tile_width>>; + using o_tile = st_bf::qo_height, fwd_attend_ker_tile_dims::tile_width>; + + using q_gl = gl; + using k_gl = gl; + using v_gl = gl; + using l_gl = gl; + using o_gl = gl; + + q_gl q; + k_gl k; + v_gl v; + l_gl l; + o_gl o; + + const int N; + const int hr; + const int max_kv_blocks_per_q; + + int32_t *__restrict__ q2k_block_sparse_index; + int32_t *__restrict__ q2k_block_sparse_num; + int32_t *__restrict__ block_size; +}; + + +template +__global__ __launch_bounds__(128, 4) +void fwd_attend_ker(const __grid_constant__ fwd_globals g) { // use block size of 64 + extern __shared__ int __shm[]; + tma_swizzle_allocator al((int*)&__shm[0]); + + using K = fwd_attend_ker_tile_dims; + + using q_tile = st_bf<64, K::tile_width>; + using k_tile = st_bf<64, K::tile_width>; + using v_tile = st_bf<64, K::tile_width>; + using l_col_vec = col_vec>; + using o_tile = st_bf<64, K::tile_width>; + + q_tile (&q_smem)[1] = al.allocate(); + + k_tile (&k_smem)[1] = al.allocate(); + + v_tile (&v_smem)[1] = al.allocate(); + + l_col_vec (&l_smem)[1] = al.allocate(); + + auto (*o_smem) = reinterpret_cast(q_smem); + + int kv_head_idx = blockIdx.y / g.hr; + int seq_idx = blockIdx.x; + + int32_t* q2k_block_sparse_index_ptr = g.q2k_block_sparse_index + blockIdx.z * gridDim.y * gridDim.x * g.max_kv_blocks_per_q + blockIdx.y * gridDim.x * g.max_kv_blocks_per_q + blockIdx.x * g.max_kv_blocks_per_q; + int32_t* q2k_block_sparse_num_ptr = g.q2k_block_sparse_num + blockIdx.z * gridDim.y * gridDim.x + blockIdx.y * gridDim.x + blockIdx.x; + int32_t kv_blocks = q2k_block_sparse_num_ptr[0]; + __shared__ kittens::semaphore qsmem_semaphore, k_smem_arrived, v_smem_arrived; + if (threadIdx.x == 0) { + int32_t kv_block_index = q2k_block_sparse_index_ptr[0]; + + init_semaphore(qsmem_semaphore, 0, 1); + init_semaphore(k_smem_arrived, 0, 1); + init_semaphore(v_smem_arrived, 0, 1); + + // preload q block + coord q_tile_idx = {blockIdx.z, blockIdx.y, seq_idx, 0}; + tma::expect_bytes(qsmem_semaphore, sizeof(q_smem)); + tma::load_async(q_smem[0], g.q, q_tile_idx, qsmem_semaphore); + + // preload the zeroth block of kv + tma::expect_bytes(k_smem_arrived, sizeof(k_tile)); + coord k_tile_idx = {blockIdx.z, kv_head_idx, kv_block_index, 0}; + tma::load_async(k_smem[0], g.k, k_tile_idx, k_smem_arrived); + + tma::expect_bytes(v_smem_arrived, sizeof(v_tile)); + coord v_tile_idx = {blockIdx.z, kv_head_idx, kv_block_index, 0}; + tma::load_async(v_smem[0], g.v, v_tile_idx, v_smem_arrived); + } + __syncthreads(); + + rt_fl<16, 64> att_block; + rt_bf<16, 64> att_block_mma; + rt_fl<16, K::tile_width> o_reg; + + col_vec> max_vec, norm_vec, max_vec_last_scaled, max_vec_scaled; + + neg_infty(max_vec); + zero(norm_vec); + zero(o_reg); + + // wait for q block + wait(qsmem_semaphore, 0); + + for (int kv_idx = 0; kv_idx < kv_blocks - 1; kv_idx++) { + // preload kv index + int32_t kv_block_index = q2k_block_sparse_index_ptr[kv_idx + 1]; + + // wait k + wait(k_smem_arrived, kv_idx % 2); + + // compute QK^T + warpgroup::mm_ABt(att_block, q_smem[0], k_smem[0]); + + copy(max_vec_last_scaled, max_vec); + if constexpr (D == 64) { mul(max_vec_last_scaled, max_vec_last_scaled, 1.44269504089f*0.125f); } + else { mul(max_vec_last_scaled, max_vec_last_scaled, 1.44269504089f*0.08838834764f); } + + warpgroup::mma_async_wait(); + + // load K + if (threadIdx.x == 0) { + tma::expect_bytes(k_smem_arrived, sizeof(k_tile)); + coord k_tile_idx = {blockIdx.z, kv_head_idx, kv_block_index, 0}; + tma::load_async(k_smem[0], g.k, k_tile_idx, k_smem_arrived); + } + + // exp + right_fill(att_block, att_block, g.block_size[q2k_block_sparse_index_ptr[kv_idx]], base_types::constants::neg_infty()); + row_max(max_vec, att_block, max_vec); + + if constexpr (D == 64) { + mul(att_block, att_block, 1.44269504089f*0.125f); + mul(max_vec_scaled, max_vec, 1.44269504089f*0.125f); + } + else { + mul(att_block, att_block, 1.44269504089f*0.08838834764f); + mul(max_vec_scaled, max_vec, 1.44269504089f*0.08838834764f); + } + + sub_row(att_block, att_block, max_vec_scaled); + exp2(att_block, att_block); + sub(max_vec_last_scaled, max_vec_last_scaled, max_vec_scaled); + exp2(max_vec_last_scaled, max_vec_last_scaled); + mul(norm_vec, norm_vec, max_vec_last_scaled); + row_sum(norm_vec, att_block, norm_vec); + add(att_block, att_block, 0.f); + copy(att_block_mma, att_block); + mul_row(o_reg, o_reg, max_vec_last_scaled); + + // wait v + wait(v_smem_arrived, kv_idx % 2); + + // compute SV + warpgroup::mma_AB(o_reg, att_block_mma, v_smem[0]); + warpgroup::mma_async_wait(); + + // load V + if (threadIdx.x == 0) { + tma::expect_bytes(v_smem_arrived, sizeof(v_tile)); + coord v_tile_idx = {blockIdx.z, kv_head_idx, kv_block_index, 0}; + tma::load_async(v_smem[0], g.v, v_tile_idx, v_smem_arrived); + } + } + + // last iter + { + int kv_idx = kv_blocks - 1; + // wait k + wait(k_smem_arrived, kv_idx % 2); + + // compute QK^T + warpgroup::mm_ABt(att_block, q_smem[0], k_smem[0]); + + copy(max_vec_last_scaled, max_vec); + if constexpr (D == 64) { mul(max_vec_last_scaled, max_vec_last_scaled, 1.44269504089f*0.125f); } + else { mul(max_vec_last_scaled, max_vec_last_scaled, 1.44269504089f*0.08838834764f); } + + warpgroup::mma_async_wait(); + + // exp + right_fill(att_block, att_block, g.block_size[q2k_block_sparse_index_ptr[kv_idx]], base_types::constants::neg_infty()); + + row_max(max_vec, att_block, max_vec); + + if constexpr (D == 64) { + mul(att_block, att_block, 1.44269504089f*0.125f); + mul(max_vec_scaled, max_vec, 1.44269504089f*0.125f); + } + else { + mul(att_block, att_block, 1.44269504089f*0.08838834764f); + mul(max_vec_scaled, max_vec, 1.44269504089f*0.08838834764f); + } + + sub_row(att_block, att_block, max_vec_scaled); + exp2(att_block, att_block); + sub(max_vec_last_scaled, max_vec_last_scaled, max_vec_scaled); + exp2(max_vec_last_scaled, max_vec_last_scaled); + mul(norm_vec, norm_vec, max_vec_last_scaled); + row_sum(norm_vec, att_block, norm_vec); + add(att_block, att_block, 0.f); + copy(att_block_mma, att_block); + mul_row(o_reg, o_reg, max_vec_last_scaled); + + // wait v + wait(v_smem_arrived, kv_idx % 2); + + // compute SV + warpgroup::mma_AB(o_reg, att_block_mma, v_smem[0]); + warpgroup::mma_async_wait(); + } + + div_row(o_reg, o_reg, norm_vec); + warpgroup::store(o_smem[0], o_reg); + __syncthreads(); + + // TK store_async internally calls syncwarp so we need to route on warp level + if (threadIdx.x / 32 == 0) { + coord o_tile_idx = {blockIdx.z, blockIdx.y, seq_idx, 0}; + tma::store_async(g.o, o_smem[0], o_tile_idx); + } + + mul(max_vec_scaled, max_vec_scaled, 0.69314718056f); + log(norm_vec, norm_vec); + add(norm_vec, norm_vec, max_vec_scaled); + + if constexpr (D == 64) { mul(norm_vec, norm_vec, -8.0f); } + else { mul(norm_vec, norm_vec, -11.313708499f); } + + warpgroup::store(l_smem[0], norm_vec); + __syncthreads(); + + if (threadIdx.x / 32 == 0) { + coord tile_idx = {blockIdx.z, blockIdx.y, 0, seq_idx}; + tma::store_async(g.l, l_smem[0], tile_idx); + } + tma::store_async_wait(); +} + +// --------------------------------------------------------------------------------------------------- +// ----------------------------------- Backward preparation kernel ----------------------------------- +// --------------------------------------------------------------------------------------------------- + +template +struct bwd_prep_globals { + using og_tile = st_bf<4*16, D>; + using o_tile = st_bf<4*16, D>; + using d_tile = col_vec>; + + using og_gl = gl; + using o_gl = gl; + using d_gl = gl; + + og_gl og; + o_gl o; + d_gl d; +}; + +constexpr int PREP_NUM_WARPS = (1); +template +__global__ __launch_bounds__(PREP_NUM_WARPS*kittens::WARP_THREADS, (D == 64) ? 6 / PREP_NUM_WARPS : 3 / PREP_NUM_WARPS) +void bwd_attend_prep_ker(const __grid_constant__ bwd_prep_globals g) { + extern __shared__ int __shm[]; + tma_swizzle_allocator al((int*)&__shm[0]); + + int warpid = kittens::warpid(); + + using og_tile = st_bf<4*16, D>; + using o_tile = st_bf<4*16, D>; + using d_tile = col_vec>; + + og_tile (&og_smem)[PREP_NUM_WARPS] = al.allocate(); + o_tile (&o_smem) [PREP_NUM_WARPS] = al.allocate(); + d_tile (&d_smem) [PREP_NUM_WARPS] = al.allocate(); + + rt_fl<4*16, D> og_reg, o_reg; + col_vec> d_reg; + + __shared__ kittens::semaphore smem_semaphore; + + if (threadIdx.x == 0) { + init_semaphore(smem_semaphore, 0, 1); + tma::expect_bytes(smem_semaphore, sizeof(og_smem[0]) * PREP_NUM_WARPS * 2); + } + __syncthreads(); + + if (warpid == 0) { + for (int w = 0; w < PREP_NUM_WARPS; w++) { + coord tile_idx = {blockIdx.z, blockIdx.y, (blockIdx.x * PREP_NUM_WARPS) + w, 0}; + tma::load_async(o_smem[w], g.o, tile_idx, smem_semaphore); + tma::load_async(og_smem[w], g.og, tile_idx, smem_semaphore); + } + } + + wait(smem_semaphore, 0); + load(o_reg, o_smem[warpid]); + load(og_reg, og_smem[warpid]); + mul(og_reg, og_reg, o_reg); + row_sum(d_reg, og_reg); + store(d_smem[warpid], d_reg); + __syncthreads(); + + if (warpid == 0) { + for (int w = 0; w < PREP_NUM_WARPS; w++) { + coord tile_idx = {blockIdx.z, blockIdx.y, 0, (blockIdx.x * PREP_NUM_WARPS) + w}; + tma::store_async(g.d, d_smem[w], tile_idx); + } + } + tma::store_async_wait(); +} + +template struct bwd_attend_ker_tile_dims {}; +template<> struct bwd_attend_ker_tile_dims<64> { + constexpr static int tile_width = (64); + constexpr static int tile_h = (4*16); + constexpr static int tile_h_qo = (4*16); +}; +template<> struct bwd_attend_ker_tile_dims<128> { + constexpr static int tile_width = (128); + constexpr static int tile_h = (4*16); + constexpr static int tile_h_qo = (4*16); +}; + +template +struct bwd_globals { + using G = bwd_attend_ker_tile_dims; + + using q_tile = st_bf; + using k_tile = st_bf; + using v_tile = st_bf; + using og_tile = st_bf; + using qg_tile = st_fl; + using kg_tile = st_fl; + using vg_tile = st_fl; + using l_tile = row_vec>; + using d_tile = row_vec>; + + using q_gl = gl; + using k_gl = gl; + using v_gl = gl; + + using og_gl = gl; + + using qg_gl = gl; + using kg_gl = gl; + using vg_gl = gl; + + using l_gl = gl; + using d_gl = gl; + + q_gl q; + k_gl k; + v_gl v; + og_gl og; + qg_gl qg; + kg_gl kg; + vg_gl vg; + l_gl l; + d_gl d; + + const int N; + const int hr; + const int max_q_blocks_per_kv; + + int32_t *__restrict__ k2q_block_sparse_index; + int32_t *__restrict__ k2q_block_sparse_num; + int32_t *__restrict__ block_size; +}; + +__device__ static inline void +stream_tile(auto ®_tile, auto &smem_vec, int tic) { + #pragma unroll + for(int i = 0; i < 4; i++) { + int base_col = 16*i + 2*(kittens::laneid()%4); + reg_tile.tiles[0][i].data[0] = *(float2*)&smem_vec[tic][base_col + 0]; + reg_tile.tiles[0][i].data[1] = *(float2*)&smem_vec[tic][base_col + 0]; + reg_tile.tiles[0][i].data[2] = *(float2*)&smem_vec[tic][base_col + 8]; + reg_tile.tiles[0][i].data[3] = *(float2*)&smem_vec[tic][base_col + 8]; + } +} + +__device__ static inline void +stream_sub_tile(auto ®_tile, auto &smem_vec, int tic) { + #pragma unroll + for(int i = 0; i < 4; i++) { + int base_col = 16*i + 2*(laneid()%4); + reg_tile.tiles[0][i].data[0] = base_ops::sub::template op(reg_tile.tiles[0][i].data[0], *(float2*)&smem_vec[tic][base_col + 0]); + reg_tile.tiles[0][i].data[1] = base_ops::sub::template op(reg_tile.tiles[0][i].data[1], *(float2*)&smem_vec[tic][base_col + 0]); + reg_tile.tiles[0][i].data[2] = base_ops::sub::template op(reg_tile.tiles[0][i].data[2], *(float2*)&smem_vec[tic][base_col + 8]); + reg_tile.tiles[0][i].data[3] = base_ops::sub::template op(reg_tile.tiles[0][i].data[3], *(float2*)&smem_vec[tic][base_col + 8]); + } +} + +template +__global__ __launch_bounds__(128, (D == 64)? 3 : 2) +void bwd_attend_ker(const __grid_constant__ bwd_globals g) { + extern __shared__ int __shm[]; + tma_swizzle_allocator al((int*)&__shm[0]); + + const int N = g.N, hr = g.hr; + using G = bwd_attend_ker_tile_dims; + + using kg_tile = st_fl; + using vg_tile = st_fl; + using k_tile = st_bf; + using v_tile = st_bf; + using q_tile = st_bf; + using og_tile = st_bf; + using qg_tile = st_fl; + using l_tile = row_vec>; + using d_tile = row_vec>; + using attn_tile = st_bf; + + k_tile (&k_smem) [1] = al.allocate(); + v_tile (&v_smem) [1] = al.allocate(); + + q_tile (&q_smem) [1] = al.allocate(); + og_tile (&og_smem)[1] = al.allocate(); + qg_tile (&qg_smem) = al.allocate(); + + l_tile (&l_smem)[1] = al.allocate(); + d_tile (&d_smem)[1] = al.allocate(); + kg_tile (*kg_smem) = reinterpret_cast(&k_smem[0].data[0]); + vg_tile (*vg_smem) = reinterpret_cast(&q_smem[0].data[0]); + + attn_tile (&ds_smem_t)[1] = al.allocate(); + + const int warpid = kittens::warpid(); + const int warpgroupid = warpid/kittens::WARPGROUP_WARPS; + const int kv_head_idx = (blockIdx.y) / hr; + + + int32_t *__restrict__ k2q_block_sparse_index_ptr = g.k2q_block_sparse_index + blockIdx.z * gridDim.y * gridDim.x * g.max_q_blocks_per_kv + blockIdx.y * gridDim.x * g.max_q_blocks_per_kv + blockIdx.x * g.max_q_blocks_per_kv; + int32_t *__restrict__ k2q_block_sparse_num_ptr = g.k2q_block_sparse_num + blockIdx.z * gridDim.y * gridDim.x + blockIdx.y * gridDim.x + blockIdx.x; + const int qo_blocks = *k2q_block_sparse_num_ptr; + + if (qo_blocks <= 0) { + return; + } + + __shared__ kittens::semaphore kv_b, q_b[1], o_b[1], vec_b[1]; + + int32_t store_qg_block_index; + int32_t load_q_block_index; + + if (threadIdx.x == 0) { + load_q_block_index = k2q_block_sparse_index_ptr[0]; + + init_semaphore(kv_b, 0, 1); + + init_semaphore(q_b[0], 0, 1); + init_semaphore(o_b[0], 0, 1); + init_semaphore(vec_b[0], 0, 1); + + // preload KV + tma::expect_bytes(kv_b, sizeof(k_smem[0]) + sizeof(v_smem[0])); + coord tile_idx_kv = {blockIdx.z, kv_head_idx, blockIdx.x, 0}; + tma::load_async(k_smem[0], g.k, tile_idx_kv, kv_b); + tma::load_async(v_smem[0], g.v, tile_idx_kv, kv_b); + + // preload og, vec and q + coord tile_idx_qo = {blockIdx.z, blockIdx.y, load_q_block_index, 0}; + coord vec_idx = {blockIdx.z, blockIdx.y, 0, load_q_block_index}; + + tma::expect_bytes(o_b[0], sizeof(og_smem[0])); + tma::load_async(og_smem[0], g.og, tile_idx_qo, o_b[0]); + + tma::expect_bytes(vec_b[0], sizeof(l_smem[0]) + sizeof(d_smem[0])); + tma::load_async(l_smem[0], g.l, vec_idx, vec_b[0]); + tma::load_async(d_smem[0], g.d, vec_idx, vec_b[0]); + + tma::expect_bytes(q_b[0], sizeof(q_smem[0])); + tma::load_async(q_smem[0], g.q, tile_idx_qo, q_b[0]); + } + __syncthreads(); + + rt_fl<16, G::tile_width> kg_reg, vg_reg; + + row_vec> row_reg; + + rt_fl<16, 64> s_block_t, p_block_t; + rt_fl<16, 64> ds_block_t, dp_block_t; + rt_bf<16, 64> ds_block_t_mma, p_block_t_mma; + + zero(kg_reg); + zero(vg_reg); + + // wait for kv + wait(kv_b, 0); + int fill_start = g.block_size[blockIdx.x] - 16 * kittens::warpid(); + for (int qo_idx = 0; qo_idx < qo_blocks - 1; qo_idx++) { + // preload q index + store_qg_block_index = load_q_block_index; + load_q_block_index = k2q_block_sparse_index_ptr[qo_idx + 1]; + + wait(o_b[0], qo_idx % 2); + warpgroup::mm_ABt(dp_block_t, v_smem[0], og_smem[0]); // dP^T = VdO^T + warpgroup::mma_commit_group(); // ! do not wait + + wait(vec_b[0], qo_idx % 2); + stream_tile(s_block_t, l_smem, 0); + wait(q_b[0], qo_idx % 2); + warpgroup::mma_ABt(s_block_t, k_smem[0], q_smem[0]); // S^T = KQ^T - l + warpgroup::mma_commit_group(); + warpgroup::mma_async_wait(); + + if constexpr (D == 64) { mul(s_block_t, s_block_t, 1.44269504089f*0.125f); } + else { mul(s_block_t, s_block_t, 1.44269504089f*0.08838834764f); } + + lower_fill(s_block_t, s_block_t, fill_start, base_types::constants::neg_infty()); + exp2(s_block_t, s_block_t); // P_i + copy(p_block_t, s_block_t); + copy(p_block_t_mma, s_block_t); + stream_sub_tile(dp_block_t, d_smem, 0); // dP - D + mul(ds_block_t, p_block_t, dp_block_t); // dS = P \odot (dP - D) + + if constexpr (D == 64) { mul(ds_block_t, ds_block_t, 0.125f); } + else { mul(ds_block_t, ds_block_t, 0.08838834764f); } + + // load vec + if (threadIdx.x == 0) { + coord vec_idx = {blockIdx.z, blockIdx.y, 0, load_q_block_index}; + tma::expect_bytes(vec_b[0], sizeof(l_smem[0]) + sizeof(d_smem[0])); + tma::load_async(l_smem[0], g.l, vec_idx, vec_b[0]); + tma::load_async(d_smem[0], g.d, vec_idx, vec_b[0]); + } + + warpgroup::mma_AB(vg_reg, p_block_t_mma, og_smem[0]); // dV += P^TdO + warpgroup::mma_commit_group(); + copy(ds_block_t_mma, ds_block_t); + warpgroup::store(ds_smem_t[0], ds_block_t); + warpgroup::mma_async_wait(); + + // load og + if (threadIdx.x == 0) { + coord tile_idx = {blockIdx.z, blockIdx.y, load_q_block_index, 0}; + tma::expect_bytes(o_b[0], sizeof(og_smem[0])); + tma::load_async(og_smem[0], g.og, tile_idx, o_b[0]); + } + + warpgroup::mma_AB(kg_reg, ds_block_t_mma, q_smem[0]); // dK += dS^TQ + warpgroup::mma_commit_group(); + warpgroup::mma_async_wait(); + + // load q + if (threadIdx.x == 0) { + coord q_tile_idx = {blockIdx.z, blockIdx.y, load_q_block_index, 0}; + tma::expect_bytes(q_b[0], sizeof(q_smem[0])); + tma::load_async(q_smem[0], g.q, q_tile_idx, q_b[0]); + } + + rt_fl<16, G::tile_width> qg_reg; + __syncthreads(); // wait for sd_smem shared memory write + warpgroup::mm_AtB(qg_reg, ds_smem_t[0], k_smem[0]); //delat dQ = dSK + warpgroup::mma_commit_group(); + warpgroup::mma_async_wait(); + // store qg to shared memory + warpgroup::store(qg_smem, qg_reg); + __syncthreads(); + + // store and add dQ to global memory + if (threadIdx.x / 32 == 0) { + coord tile_idx = {blockIdx.z, blockIdx.y, store_qg_block_index, 0}; + tma::store_add_async(g.qg, qg_smem, tile_idx); + tma::store_async_wait(); + } + } + + // last iter + { + int qo_idx = qo_blocks - 1; + + store_qg_block_index = load_q_block_index; + + wait(o_b[0], qo_idx % 2); + warpgroup::mm_ABt(dp_block_t, v_smem[0], og_smem[0]); // dP = dOV^T + warpgroup::mma_commit_group(); // ! do not wait + + wait(vec_b[0], qo_idx % 2); + stream_tile(s_block_t, l_smem, 0); + wait(q_b[0], qo_idx % 2); + warpgroup::mma_ABt(s_block_t, k_smem[0], q_smem[0]); // S = QK^T - l + warpgroup::mma_commit_group(); + warpgroup::mma_async_wait(); + + if constexpr (D == 64) { mul(s_block_t, s_block_t, 1.44269504089f*0.125f); } + else { mul(s_block_t, s_block_t, 1.44269504089f*0.08838834764f); } + lower_fill(s_block_t, s_block_t, fill_start, base_types::constants::neg_infty()); + exp2(s_block_t, s_block_t); // P_i + copy(p_block_t, s_block_t); + copy(p_block_t_mma, s_block_t); + stream_sub_tile(dp_block_t, d_smem, 0); // dP - D + mul(ds_block_t, p_block_t, dp_block_t); // dS = P \odot (dP - D) + + if constexpr (D == 64) { mul(ds_block_t, ds_block_t, 0.125f); } + else { mul(ds_block_t, ds_block_t, 0.08838834764f); } + + warpgroup::mma_AB(vg_reg, p_block_t_mma, og_smem[0]); // dV += P^TdO + warpgroup::mma_commit_group(); + copy(ds_block_t_mma, ds_block_t); + warpgroup::store(ds_smem_t[0], ds_block_t); + warpgroup::mma_async_wait(); + + warpgroup::mma_AB(kg_reg, ds_block_t_mma, q_smem[0]); // dK += dS^TQ + warpgroup::mma_commit_group(); + warpgroup::mma_async_wait(); + + rt_fl<16, G::tile_width> qg_reg; + __syncthreads(); // wait for sd_smem shared memory write + warpgroup::mm_AtB(qg_reg, ds_smem_t[0], k_smem[0]); //delat dQ = dSK + warpgroup::mma_commit_group(); + warpgroup::mma_async_wait(); + // store qg to shared memory + warpgroup::store(qg_smem, qg_reg); + __syncthreads(); + + // store and add dQ to global memory + if (threadIdx.x / 32 == 0) { + coord tile_idx = {blockIdx.z, blockIdx.y, store_qg_block_index, 0}; + tma::store_add_async(g.qg, qg_smem, tile_idx); + tma::store_async_wait(); + } + } + + // store kq and vq + + // ensuring all writes are finished + __syncthreads(); + + warpgroup::store(kg_smem[0], kg_reg); + __syncthreads(); + if (threadIdx.x / 32 == 0) { + coord tile_idx = {blockIdx.z, kv_head_idx, blockIdx.x, 0}; + tma::store_add_async(g.kg, kg_smem[0], tile_idx); + tma::store_commit_group(); + } + + warpgroup::store(vg_smem[0], vg_reg); + __syncthreads(); + if (kittens::warpid() % 4 == 0) { + coord tile_idx = {blockIdx.z, kv_head_idx, blockIdx.x, 0}; + tma::store_add_async(g.vg, vg_smem[0], tile_idx); + tma::store_commit_group(); + } + tma::store_async_wait(); +} + + +template +void block_sparse_attention_forward_impl( + bf16* d_q, bf16* d_k, bf16* d_v, float* d_l, bf16* d_o, + int batch, int qo_heads, int kv_heads, int seq_len, int hr, + int max_kv_blocks_per_q, + int32_t* q2k_block_sparse_index_ptr, + int32_t* q2k_block_sparse_num_ptr, + int32_t* block_size_ptr, + cudaStream_t stream +) { + using K = fwd_attend_ker_tile_dims; + using q_tile = st_bf; + using k_tile = st_bf; + using v_tile = st_bf; + using l_col_vec = col_vec>; + using o_tile = st_bf; + + using q_global = gl; + using k_global = gl; + using v_global = gl; + using l_global = gl; + using o_global = gl; + + using globals = fwd_globals; + + q_global qg_arg{d_q, static_cast(batch), static_cast(qo_heads), static_cast(seq_len), static_cast(D)}; + k_global kg_arg{d_k, static_cast(batch), static_cast(kv_heads), static_cast(seq_len), static_cast(D)}; + v_global vg_arg{d_v, static_cast(batch), static_cast(kv_heads), static_cast(seq_len), static_cast(D)}; + l_global lg_arg{d_l, static_cast(batch), static_cast(qo_heads), 1U, static_cast(seq_len)}; + o_global og_arg{d_o, static_cast(batch), static_cast(qo_heads), static_cast(seq_len), static_cast(D)}; + + globals g{ + qg_arg, kg_arg, vg_arg, lg_arg, og_arg, + static_cast(seq_len), static_cast(hr), static_cast(max_kv_blocks_per_q), + q2k_block_sparse_index_ptr, q2k_block_sparse_num_ptr, block_size_ptr + }; + + // Shared memory size for the kernel + // 54000 bytes is calibrated for H100 shared memory constraints for these tile sizes + constexpr int mem_size = 54000; + + dim3 grid(seq_len/(64), qo_heads, batch); + + cudaFuncSetAttribute( + fwd_attend_ker, + cudaFuncAttributeMaxDynamicSharedMemorySize, + mem_size + ); + + fwd_attend_ker<<>>(g); +} + +template +void block_sparse_attention_backward_impl( + bf16* d_q, bf16* d_k, bf16* d_v, bf16* d_o, bf16* d_og, float* d_l, float* d_d, float* d_qg, float* d_kg, float* d_vg, + int batch, int qo_heads, int kv_heads, int seq_len, int hr, int max_q_blocks_per_kv, + int32_t* k2q_block_sparse_index_ptr, + int32_t* k2q_block_sparse_num_ptr, + int32_t* block_size_ptr, + cudaStream_t stream +) { + using G = bwd_attend_ker_tile_dims; + using og_tile = st_bf<4*16, D>; + using o_tile = st_bf<4*16, D>; + using d_tile = col_vec>; + + using og_global = gl; + using o_global = gl; + using d_global = gl; + + using prep_globals = bwd_prep_globals; + + constexpr int mem_size_prep = kittens::MAX_SHARED_MEMORY; + int threads_prep = PREP_NUM_WARPS * kittens::WARP_THREADS; + dim3 grid_bwd_prep(seq_len/(PREP_NUM_WARPS*kittens::TILE_ROW_DIM*4), qo_heads, batch); + + og_global prep_og_arg{d_og, static_cast(batch), static_cast(qo_heads), static_cast(seq_len), static_cast(D)}; + o_global prep_o_arg {d_o, static_cast(batch), static_cast(qo_heads), static_cast(seq_len), static_cast(D)}; + d_global prep_d_arg {d_d, static_cast(batch), static_cast(qo_heads), 1U, static_cast(seq_len)}; + + prep_globals bwd_g{prep_og_arg, prep_o_arg, prep_d_arg}; + + cudaFuncSetAttribute( + bwd_attend_prep_ker, + cudaFuncAttributeMaxDynamicSharedMemorySize, + mem_size_prep + ); + bwd_attend_prep_ker<<>>(bwd_g); + + using bwd_q_tile = st_bf; + using bwd_k_tile = st_bf; + using bwd_v_tile = st_bf; + using bwd_og_tile = st_bf; + using bwd_qg_tile = st_fl; + using bwd_kg_tile = st_fl; + using bwd_vg_tile = st_fl; + using bwd_l_tile = row_vec>; + using bwd_d_tile = row_vec>; + + using bwd_q_global = gl; + using bwd_k_global = gl; + using bwd_v_global = gl; + using bwd_og_global = gl; + using bwd_qg_global = gl; + using bwd_kg_global = gl; + using bwd_vg_global = gl; + using bwd_l_global = gl; + using bwd_d_global = gl; + + using bwd_global_args = bwd_globals; + + bwd_q_global bwd_q_arg {d_q, static_cast(batch), static_cast(qo_heads), static_cast(seq_len), static_cast(D)}; + bwd_k_global bwd_k_arg {d_k, static_cast(batch), static_cast(kv_heads), static_cast(seq_len), static_cast(D)}; + bwd_v_global bwd_v_arg {d_v, static_cast(batch), static_cast(kv_heads), static_cast(seq_len), static_cast(D)}; + bwd_og_global bwd_og_arg{d_og, static_cast(batch), static_cast(qo_heads), static_cast(seq_len), static_cast(D)}; + bwd_qg_global bwd_qg_arg{d_qg, static_cast(batch), static_cast(qo_heads), static_cast(seq_len), static_cast(D)}; + bwd_kg_global bwd_kg_arg{d_kg, static_cast(batch), static_cast(kv_heads), static_cast(seq_len), static_cast(D)}; + bwd_vg_global bwd_vg_arg{d_vg, static_cast(batch), static_cast(kv_heads), static_cast(seq_len), static_cast(D)}; + bwd_l_global bwd_l_arg {d_l, static_cast(batch), static_cast(qo_heads), 1U, static_cast(seq_len)}; + bwd_d_global bwd_d_arg {d_d, static_cast(batch), static_cast(qo_heads), 1U, static_cast(seq_len)}; + + bwd_global_args bwd_global{bwd_q_arg, bwd_k_arg, bwd_v_arg, bwd_og_arg, bwd_qg_arg, bwd_kg_arg, bwd_vg_arg, bwd_l_arg, bwd_d_arg, + static_cast(seq_len), static_cast(hr), static_cast(max_q_blocks_per_kv), + k2q_block_sparse_index_ptr, k2q_block_sparse_num_ptr, block_size_ptr}; + + dim3 grid_bwd_main(seq_len/64, qo_heads, batch); + int threads_main = 128; + // Calibrated shared memory sizes for different head dimensions + int bwd_mem_size = (D == 64) ? 72000 : 113000; + + cudaFuncSetAttribute( + bwd_attend_ker, + cudaFuncAttributeMaxDynamicSharedMemorySize, + bwd_mem_size + ); + bwd_attend_ker<<>>(bwd_global); +} + +#include "pyutils/torch_helpers.cuh" +#include +#include + +std::vector +block_sparse_attention_forward( + torch::Tensor q, + torch::Tensor k, + torch::Tensor v, + torch::Tensor q2k_block_sparse_index, + torch::Tensor q2k_block_sparse_num, + torch::Tensor block_size +) +{ + CHECK_INPUT(q); + CHECK_INPUT(k); + CHECK_INPUT(v); + + auto batch = q.size(0); + auto seq_len = q.size(2); + auto head_dim = q.size(3); + auto qo_heads = q.size(1); + auto kv_heads = k.size(1); + auto max_kv_blocks_per_q = q2k_block_sparse_index.size(3); + auto num_q_blocks = block_size.size(0); + TORCH_CHECK(batch==1, "Batch size dim will be removed in the future, please set batch to 1"); + TORCH_CHECK(num_q_blocks * 64 == seq_len, "This kernel supports variable block size, but it assumes the input sequence is properly padded."); + TORCH_CHECK(num_q_blocks == q2k_block_sparse_index.size(2), "Number of Q blocks does not match between q2k_block_sparse_index and block_size"); + // check to see that these dimensions match for all inputs + TORCH_CHECK(q.size(0) == batch, "Q batch dimension - idx 0 - must match for all inputs"); + TORCH_CHECK(k.size(0) == batch, "K batch dimension - idx 0 - must match for all inputs"); + TORCH_CHECK(v.size(0) == batch, "V batch dimension - idx 0 - must match for all inputs"); + TORCH_CHECK(q2k_block_sparse_index.size(0) == batch, "q2k_block_sparse_index batch dimension - idx 0 - must match for all inputs"); + TORCH_CHECK(q2k_block_sparse_num.size(0) == batch, "q2k_block_sparse_num batch dimension - idx 0 - must match for all inputs"); + + TORCH_CHECK(q.size(2) == seq_len, "Q sequence length dimension - idx 2 - must match for all inputs"); + TORCH_CHECK(k.size(2) == seq_len, "K sequence length dimension - idx 2 - must match for all inputs"); + TORCH_CHECK(v.size(2) == seq_len, "V sequence length dimension - idx 2 - must match for all inputs"); + TORCH_CHECK(q2k_block_sparse_index.size(2) == seq_len / BLOCK_M, "q2k_block_sparse_index idx 2 - must match seq_len / BLOCK_M"); + TORCH_CHECK(q2k_block_sparse_num.size(2) == seq_len / BLOCK_M, "q2k_block_sparse_num idx 2 - must match seq_len / BLOCK_M"); + + TORCH_CHECK(q.size(3) == head_dim, "Q head dimension - idx 3 - must match for all non-vector inputs"); + TORCH_CHECK(k.size(3) == head_dim, "K head dimension - idx 3 - must match for all non-vector inputs"); + TORCH_CHECK(v.size(3) == head_dim, "V head dimension - idx 3 - must match for all non-vector inputs"); + + TORCH_CHECK(qo_heads >= kv_heads, "QO heads must be greater than or equal to KV heads"); + TORCH_CHECK(qo_heads % kv_heads == 0, "QO heads must be divisible by KV heads"); + TORCH_CHECK(q.size(1) == qo_heads, "QO head dimension - idx 1 - must match for all inputs"); + TORCH_CHECK(k.size(1) == kv_heads, "KV head dimension - idx 1 - must match for all inputs"); + TORCH_CHECK(v.size(1) == kv_heads, "KV head dimension - idx 1 - must match for all inputs"); + TORCH_CHECK(q2k_block_sparse_index.size(1) == qo_heads, "q2k_block_sparse_index head dimension - idx 1 - must match for all inputs"); + TORCH_CHECK(q2k_block_sparse_num.size(1) == qo_heads, "q2k_block_sparse_num head dimension - idx 1 - must match for all inputs"); + auto hr = qo_heads / kv_heads; + + c10::BFloat16* q_ptr = q.data_ptr(); + c10::BFloat16* k_ptr = k.data_ptr(); + c10::BFloat16* v_ptr = v.data_ptr(); + + bf16* d_q = reinterpret_cast(q_ptr); + bf16* d_k = reinterpret_cast(k_ptr); + bf16* d_v = reinterpret_cast(v_ptr); + + // for the returned outputs + torch::Tensor o = torch::empty({static_cast(batch), + static_cast(qo_heads), + static_cast(seq_len), + static_cast(head_dim)}, v.options()); + + torch::Tensor l_vec = torch::empty({static_cast(batch), + static_cast(qo_heads), + static_cast(seq_len), + static_cast(1)}, + torch::TensorOptions().dtype(torch::kFloat).device(q.device()).memory_format(at::MemoryFormat::Contiguous)); + + + bf16* o_ptr = reinterpret_cast(o.data_ptr()); + bf16* d_o = reinterpret_cast(o_ptr); + + float* l_ptr = reinterpret_cast(l_vec.data_ptr()); + float* d_l = reinterpret_cast(l_ptr); + + const c10::cuda::OptionalCUDAGuard device_guard(q.device()); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); + + // Temporated implementation to avoid code duplication between head_dim=64 and 128 + if (head_dim == 64) { + block_sparse_attention_forward_impl<64>( + d_q, d_k, d_v, d_l, d_o, + batch, qo_heads, kv_heads, seq_len, hr, + max_kv_blocks_per_q, + reinterpret_cast(q2k_block_sparse_index.data_ptr()), + reinterpret_cast(q2k_block_sparse_num.data_ptr()), + reinterpret_cast(block_size.data_ptr()), + stream + ); + } else if (head_dim == 128) { + block_sparse_attention_forward_impl<128>( + d_q, d_k, d_v, d_l, d_o, + batch, qo_heads, kv_heads, seq_len, hr, + max_kv_blocks_per_q, + reinterpret_cast(q2k_block_sparse_index.data_ptr()), + reinterpret_cast(q2k_block_sparse_num.data_ptr()), + reinterpret_cast(block_size.data_ptr()), + stream + ); + } else { + TORCH_CHECK(false, "Unsupported head_dim: ", head_dim, ". Only 64 and 128 are supported."); + } + + return {o, l_vec}; + //cudadevicesynchronize(); +} + +std::vector +block_sparse_attention_backward(torch::Tensor q, + torch::Tensor k, + torch::Tensor v, + torch::Tensor o, + torch::Tensor l_vec, + torch::Tensor og, + torch::Tensor k2q_block_sparse_index, + torch::Tensor k2q_block_sparse_num, + torch::Tensor block_size) +{ + CHECK_INPUT(q); + CHECK_INPUT(k); + CHECK_INPUT(v); + CHECK_INPUT(l_vec); + CHECK_INPUT(o); + CHECK_INPUT(og); + + auto batch = q.size(0); + auto seq_len = q.size(2); + auto head_dim = q.size(3); + auto max_q_blocks_per_kv = k2q_block_sparse_index.size(3); + TORCH_CHECK(k2q_block_sparse_index.size(2) == block_size.size(0), "k2q_block_sparse_index.size(2) must match block_size.size(0)"); + // check to see that these dimensions match for all inputs + TORCH_CHECK(q.size(0) == batch, "Q batch dimension - idx 0 - must match for all inputs"); + TORCH_CHECK(k.size(0) == batch, "K batch dimension - idx 0 - must match for all inputs"); + TORCH_CHECK(v.size(0) == batch, "V batch dimension - idx 0 - must match for all inputs"); + TORCH_CHECK(l_vec.size(0) == batch, "L batch dimension - idx 0 - must match for all inputs"); + TORCH_CHECK(o.size(0) == batch, "O batch dimension - idx 0 - must match for all inputs"); + TORCH_CHECK(og.size(0) == batch, "OG batch dimension - idx 0 - must match for all inputs"); + TORCH_CHECK(k2q_block_sparse_index.size(0) == batch, "k2q_block_sparse_index batch dimension - idx 0 - must match for all inputs"); + TORCH_CHECK(k2q_block_sparse_num.size(0) == batch, "k2q_block_sparse_num batch dimension - idx 0 - must match for all inputs"); + + + TORCH_CHECK(q.size(2) == seq_len, "Q sequence length dimension - idx 2 - must match for all inputs"); + TORCH_CHECK(k.size(2) == seq_len, "K sequence length dimension - idx 2 - must match for all inputs"); + TORCH_CHECK(v.size(2) == seq_len, "V sequence length dimension - idx 2 - must match for all inputs"); + TORCH_CHECK(l_vec.size(2) == seq_len, "L sequence length dimension - idx 2 - must match for all inputs"); + TORCH_CHECK(o.size(2) == seq_len, "O sequence length dimension - idx 2 - must match for all inputs"); + TORCH_CHECK(og.size(2) == seq_len, "OG sequence length dimension - idx 2 - must match for all inputs"); + TORCH_CHECK(k2q_block_sparse_index.size(2) == seq_len / BLOCK_N, "k2q_block_sparse_index idx 2 - must match seq_len / BLOCK_N"); + TORCH_CHECK(k2q_block_sparse_num.size(2) == seq_len / BLOCK_N, "k2q_block_sparse_num idx 2 - must match seq_len / BLOCK_N"); + + TORCH_CHECK(q.size(3) == head_dim, "Q head dimension - idx 3 - must match for all non-vector inputs"); + TORCH_CHECK(k.size(3) == head_dim, "K head dimension - idx 3 - must match for all non-vector inputs"); + TORCH_CHECK(v.size(3) == head_dim, "V head dimension - idx 3 - must match for all non-vector inputs"); + TORCH_CHECK(o.size(3) == head_dim, "O head dimension - idx 3 - must match for all non-vector inputs"); + TORCH_CHECK(og.size(3) == head_dim, "OG head dimension - idx 3 - must match for all non-vector inputs"); + + + + auto qo_heads = q.size(1); + auto kv_heads = k.size(1); + + TORCH_CHECK(qo_heads >= kv_heads, "Q heads must be greater than or equal to K and V heads"); + TORCH_CHECK(qo_heads % kv_heads == 0, "Q heads must be divisible by KV heads"); + + TORCH_CHECK(q.size(1) == qo_heads, "Q heads dimension - idx 1 - must match for all inputs"); + TORCH_CHECK(l_vec.size(1) == qo_heads, "L heads dimension - idx 1 - must match for all inputs"); + TORCH_CHECK(o.size(1) == qo_heads, "O heads dimension - idx 1 - must match for all inputs"); + TORCH_CHECK(og.size(1) == qo_heads, "OG heads dimension - idx 1 - must match for all inputs"); + TORCH_CHECK(k.size(1) == kv_heads, "K heads dimension - idx 1 - must match for all inputs"); + TORCH_CHECK(v.size(1) == kv_heads, "V heads dimension - idx 1 - must match for all inputs"); + TORCH_CHECK(k2q_block_sparse_index.size(1) == kv_heads, "k2q_block_sparse_index heads dimension - idx 1 - must match for all inputs"); + TORCH_CHECK(k2q_block_sparse_num.size(1) == kv_heads, "k2q_block_sparse_num heads dimension - idx 1 - must match for all inputs"); + auto hr = qo_heads / kv_heads; + + c10::BFloat16* q_ptr = q.data_ptr(); + c10::BFloat16* k_ptr = k.data_ptr(); + c10::BFloat16* v_ptr = v.data_ptr(); + c10::BFloat16* o_ptr = o.data_ptr(); + c10::BFloat16* og_ptr = og.data_ptr(); + float* l_ptr = l_vec.data_ptr(); + + torch::Tensor qg = torch::zeros({static_cast(batch), + static_cast(qo_heads), + static_cast(seq_len), + static_cast(head_dim)}, l_vec.options()); + torch::Tensor kg = torch::zeros({static_cast(batch), + static_cast(kv_heads), + static_cast(seq_len), + static_cast(head_dim)}, l_vec.options()); + torch::Tensor vg = torch::zeros({static_cast(batch), + static_cast(kv_heads), + static_cast(seq_len), + static_cast(head_dim)}, l_vec.options()); + + torch::Tensor d_vec = torch::empty({static_cast(batch), + static_cast(qo_heads), + static_cast(seq_len), + static_cast(1)}, l_vec.options()); + + float* qg_ptr = qg.data_ptr(); + float* kg_ptr = kg.data_ptr(); + float* vg_ptr = vg.data_ptr(); + float* d_ptr = d_vec.data_ptr(); + + bf16* d_q = reinterpret_cast(q_ptr); + bf16* d_k = reinterpret_cast(k_ptr); + bf16* d_v = reinterpret_cast(v_ptr); + bf16* d_o = reinterpret_cast(o_ptr); + bf16* d_og = reinterpret_cast(og_ptr); + float* d_l = reinterpret_cast(l_ptr); + float* d_d = reinterpret_cast(d_ptr); + float* d_qg = reinterpret_cast(qg_ptr); + float* d_kg = reinterpret_cast(kg_ptr); + float* d_vg = reinterpret_cast(vg_ptr); + + constexpr int mem_size = kittens::MAX_SHARED_MEMORY; + int threads = PREP_NUM_WARPS * kittens::WARP_THREADS; + + //cudadevicesynchronize(); + const c10::cuda::OptionalCUDAGuard device_guard(q.device()); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); + + // cudaStreamSynchronize(stream); + + // TORCH_CHECK(seq_len % (4*kittens::TILE_DIM*4) == 0, "sequence length must be divisible by 256"); + dim3 grid_bwd(seq_len/(PREP_NUM_WARPS*kittens::TILE_ROW_DIM*4), qo_heads, batch); + + if (head_dim == 64) { + using og_tile = st_bf<4*16, 64>; + using o_tile = st_bf<4*16, 64>; + using d_tile = col_vec>; + + using og_global = gl; + using o_global = gl; + using d_global = gl; + + using bwd_prep_globals = bwd_prep_globals<64>; + + og_global prep_og_arg{d_og, static_cast(batch), static_cast(qo_heads), static_cast(seq_len), 64U}; + o_global prep_o_arg {d_o, static_cast(batch), static_cast(qo_heads), static_cast(seq_len), 64U}; + d_global prep_d_arg {d_d, static_cast(batch), static_cast(qo_heads), 1U, static_cast(seq_len)}; + + bwd_prep_globals bwd_g{prep_og_arg, prep_o_arg, prep_d_arg}; + + cudaFuncSetAttribute( + bwd_attend_prep_ker<64>, + cudaFuncAttributeMaxDynamicSharedMemorySize, + mem_size + ); + + bwd_attend_prep_ker<64><<>>(bwd_g); + + using bwd_q_tile = st_bf::tile_h_qo, bwd_attend_ker_tile_dims<64>::tile_width>; + using bwd_k_tile = st_bf::tile_h, bwd_attend_ker_tile_dims<64>::tile_width>; + using bwd_v_tile = st_bf::tile_h, bwd_attend_ker_tile_dims<64>::tile_width>; + using bwd_og_tile = st_bf::tile_h_qo, bwd_attend_ker_tile_dims<64>::tile_width>; + using bwd_qg_tile = st_fl::tile_h_qo, bwd_attend_ker_tile_dims<64>::tile_width>; + using bwd_kg_tile = st_fl::tile_h, bwd_attend_ker_tile_dims<64>::tile_width>; + using bwd_vg_tile = st_fl::tile_h, bwd_attend_ker_tile_dims<64>::tile_width>; + using bwd_l_tile = row_vec::tile_h_qo, bwd_attend_ker_tile_dims<64>::tile_h>>; + using bwd_d_tile = row_vec::tile_h_qo, bwd_attend_ker_tile_dims<64>::tile_h>>; + + using bwd_q_global = gl; + using bwd_k_global = gl; + using bwd_v_global = gl; + + using bwd_og_global = gl; + + using bwd_qg_global = gl; + using bwd_kg_global = gl; + using bwd_vg_global = gl; + + using bwd_l_global = gl; + using bwd_d_global = gl; + + using bwd_global_args = bwd_globals<64>; + + bwd_q_global bwd_q_arg {d_q, static_cast(batch), static_cast(qo_heads), static_cast(seq_len), 64U}; + bwd_k_global bwd_k_arg {d_k, static_cast(batch), static_cast(kv_heads), static_cast(seq_len), 64U}; + bwd_v_global bwd_v_arg {d_v, static_cast(batch), static_cast(kv_heads), static_cast(seq_len), 64U}; + bwd_og_global bwd_og_arg{d_og, static_cast(batch), static_cast(qo_heads), static_cast(seq_len), 64U}; + bwd_qg_global bwd_qg_arg{d_qg, static_cast(batch), static_cast(qo_heads), static_cast(seq_len), 64U}; + bwd_kg_global bwd_kg_arg{d_kg, static_cast(batch), static_cast(kv_heads), static_cast(seq_len), 64U}; + bwd_vg_global bwd_vg_arg{d_vg, static_cast(batch), static_cast(kv_heads), static_cast(seq_len), 64U}; + bwd_l_global bwd_l_arg {d_l, static_cast(batch), static_cast(qo_heads), 1U, static_cast(seq_len)}; + bwd_d_global bwd_d_arg {d_d, static_cast(batch), static_cast(qo_heads), 1U, static_cast(seq_len)}; + + bwd_global_args bwd_global{bwd_q_arg, + bwd_k_arg, + bwd_v_arg, + bwd_og_arg, + bwd_qg_arg, + bwd_kg_arg, + bwd_vg_arg, + bwd_l_arg, + bwd_d_arg, + static_cast(seq_len), + static_cast(hr), + static_cast(max_q_blocks_per_kv), + reinterpret_cast(k2q_block_sparse_index.data_ptr()), + reinterpret_cast(k2q_block_sparse_num.data_ptr()), + reinterpret_cast(block_size.data_ptr())}; + + dim3 grid_bwd_2(seq_len/64, qo_heads, batch); + threads = 128; + + //cudadevicesynchronize(); + + { + cudaFuncSetAttribute( + bwd_attend_ker<64>, + cudaFuncAttributeMaxDynamicSharedMemorySize, + 72000 + ); + // cudaFuncSetAttribute( + // bwd_attend_ker<64>, + // cudaFuncAttributePreferredSharedMemoryCarveout, + // 85 + // ); + + bwd_attend_ker<64><<>>(bwd_global); + } + + // CHECK_CUDA_ERROR(cudaGetLastError()); + // cudaStreamSynchronize(stream); + //cudadevicesynchronize(); + // const auto kernel_end = std::chrono::high_resolution_clock::now(); + // std::cout << "Kernel Time: " << std::chrono::duration_cast(kernel_end - start).count() << "us" << std::endl; + // std::cout << "---" << std::endl; + } + + if (head_dim == 128) { + using og_tile = st_bf<4*16, 128>; + using o_tile = st_bf<4*16, 128>; + using d_tile = col_vec>; + + using og_global = gl; + using o_global = gl; + using d_global = gl; + + using bwd_prep_globals = bwd_prep_globals<128>; + + og_global prep_og_arg{d_og, static_cast(batch), static_cast(qo_heads), static_cast(seq_len), 128U}; + o_global prep_o_arg {d_o, static_cast(batch), static_cast(qo_heads), static_cast(seq_len), 128U}; + d_global prep_d_arg {d_d, static_cast(batch), static_cast(qo_heads), 1U, static_cast(seq_len)}; + + bwd_prep_globals bwd_g{prep_og_arg, prep_o_arg, prep_d_arg}; + + cudaFuncSetAttribute( + bwd_attend_prep_ker<128>, + cudaFuncAttributeMaxDynamicSharedMemorySize, + mem_size + ); + + bwd_attend_prep_ker<128><<>>(bwd_g); + + using bwd_q_tile = st_bf::tile_h_qo, bwd_attend_ker_tile_dims<128>::tile_width>; + using bwd_k_tile = st_bf::tile_h, bwd_attend_ker_tile_dims<128>::tile_width>; + using bwd_v_tile = st_bf::tile_h, bwd_attend_ker_tile_dims<128>::tile_width>; + using bwd_og_tile = st_bf::tile_h_qo, bwd_attend_ker_tile_dims<128>::tile_width>; + using bwd_qg_tile = st_fl::tile_h_qo, bwd_attend_ker_tile_dims<128>::tile_width>; + using bwd_kg_tile = st_fl::tile_h, bwd_attend_ker_tile_dims<128>::tile_width>; + using bwd_vg_tile = st_fl::tile_h, bwd_attend_ker_tile_dims<128>::tile_width>; + using bwd_l_tile = row_vec::tile_h_qo, bwd_attend_ker_tile_dims<128>::tile_h>>; + using bwd_d_tile = row_vec::tile_h_qo, bwd_attend_ker_tile_dims<128>::tile_h>>; + + using bwd_q_global = gl; + using bwd_k_global = gl; + using bwd_v_global = gl; + + using bwd_og_global = gl; + + using bwd_qg_global = gl; + using bwd_kg_global = gl; + using bwd_vg_global = gl; + + using bwd_l_global = gl; + using bwd_d_global = gl; + + using bwd_global_args = bwd_globals<128>; + + bwd_q_global bwd_q_arg {d_q, static_cast(batch), static_cast(qo_heads), static_cast(seq_len), 128U}; + bwd_k_global bwd_k_arg {d_k, static_cast(batch), static_cast(kv_heads), static_cast(seq_len), 128U}; + bwd_v_global bwd_v_arg {d_v, static_cast(batch), static_cast(kv_heads), static_cast(seq_len), 128U}; + bwd_og_global bwd_og_arg{d_og, static_cast(batch), static_cast(qo_heads), static_cast(seq_len), 128U}; + bwd_qg_global bwd_qg_arg{d_qg, static_cast(batch), static_cast(qo_heads), static_cast(seq_len), 128U}; + bwd_kg_global bwd_kg_arg{d_kg, static_cast(batch), static_cast(kv_heads), static_cast(seq_len), 128U}; + bwd_vg_global bwd_vg_arg{d_vg, static_cast(batch), static_cast(kv_heads), static_cast(seq_len), 128U}; + bwd_l_global bwd_l_arg {d_l, static_cast(batch), static_cast(qo_heads), 1U, static_cast(seq_len)}; + bwd_d_global bwd_d_arg {d_d, static_cast(batch), static_cast(qo_heads), 1U, static_cast(seq_len)}; + + bwd_global_args bwd_global{bwd_q_arg, + bwd_k_arg, + bwd_v_arg, + bwd_og_arg, + bwd_qg_arg, + bwd_kg_arg, + bwd_vg_arg, + bwd_l_arg, + bwd_d_arg, + static_cast(seq_len), + static_cast(hr), + static_cast(max_q_blocks_per_kv), + reinterpret_cast(k2q_block_sparse_index.data_ptr()), + reinterpret_cast(k2q_block_sparse_num.data_ptr()), + reinterpret_cast(block_size.data_ptr())}; + + dim3 grid_bwd_2(seq_len/64, qo_heads, batch); + threads = 128; + + //cudadevicesynchronize(); + + { + cudaFuncSetAttribute( + bwd_attend_ker<128>, + cudaFuncAttributeMaxDynamicSharedMemorySize, + 113000 + ); + + bwd_attend_ker<128><<>>(bwd_global); + } + + // cudaStreamSynchronize(stream); + //cudadevicesynchronize(); + } + + return {qg, kg, vg}; + //cudadevicesynchronize(); +} diff --git a/csrc/fastvideo_kernel/csrc/st_attn.cpp b/csrc/fastvideo_kernel/csrc/st_attn.cpp new file mode 100644 index 000000000..41a63c17f --- /dev/null +++ b/csrc/fastvideo_kernel/csrc/st_attn.cpp @@ -0,0 +1,23 @@ +#include +#include + +#include +#include +#include + +#include + +#ifdef TK_COMPILE_ST_ATTN +extern torch::Tensor sta_forward( + torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o, int kernel_t_size, int kernel_w_size, int kernel_h_size, int text_length, bool process_text, bool has_text, int kernel_aspect_ratio_flag +); +#endif + + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.doc() = "Sliding Block Attention Kernels"; // optional module docstring + +#ifdef TK_COMPILE_ST_ATTN + m.def("sta_fwd", torch::wrap_pybind_function(sta_forward), "sliding tile attention, assuming tile size is (6,8,8)"); +#endif +} diff --git a/csrc/fastvideo_kernel/csrc/st_attn_h100.cu b/csrc/fastvideo_kernel/csrc/st_attn_h100.cu new file mode 100644 index 000000000..0ad839155 --- /dev/null +++ b/csrc/fastvideo_kernel/csrc/st_attn_h100.cu @@ -0,0 +1,573 @@ +// # Define TORCH_COMPILE macro + +#include "kittens.cuh" +#include +#include +#include +#include + +// #define CLAMP(value, min, max) ((value) < (min) ? (min) : ((value) > (max) ? (max) : (value))) +__device__ __forceinline__ int clamp_int(int value, int min, int max) { + return (value < min) ? min : ((value > max) ? max : value); +} +// #define ABS(x) ((x) < 0 ? -(x) : (x)) +__device__ __forceinline__ int abs_int(int value) { + return (value < 0) ? -value : value; +} + + +constexpr int CONSUMER_WARPGROUPS = (3); +constexpr int PRODUCER_WARPGROUPS = (1); +constexpr int NUM_WARPGROUPS = (CONSUMER_WARPGROUPS+PRODUCER_WARPGROUPS); +constexpr int NUM_WORKERS = (NUM_WARPGROUPS*kittens::WARPGROUP_WARPS); + +using namespace kittens; +namespace cg = cooperative_groups; + +template struct fwd_attend_ker_tile_dims {}; +template<> struct fwd_attend_ker_tile_dims<64> { + constexpr static int tile_width = (64); + constexpr static int qo_height = (4*16); + constexpr static int kv_height = (8*16); + constexpr static int stages = (4); +}; +template<> struct fwd_attend_ker_tile_dims<128> { + constexpr static int tile_width = (128); + constexpr static int qo_height = (4*16); + constexpr static int kv_height = (8*16); + constexpr static int stages = (2); +}; + +template struct fwd_globals { + using q_tile = st_bf::qo_height, fwd_attend_ker_tile_dims::tile_width>; + using k_tile = st_bf::kv_height, fwd_attend_ker_tile_dims::tile_width>; + using v_tile = st_bf::kv_height, fwd_attend_ker_tile_dims::tile_width>; + using l_col_vec = col_vec::qo_height, fwd_attend_ker_tile_dims::tile_width>>; + using o_tile = st_bf::qo_height, fwd_attend_ker_tile_dims::tile_width>; + + using q_gl = gl; + using k_gl = gl; + using v_gl = gl; + using l_gl = gl; + using o_gl = gl; + + q_gl q; + k_gl k; + v_gl v; + l_gl l; + o_gl o; + + const int N; + const int text_L; + const int hr; +}; + + +template +__global__ __launch_bounds__((NUM_WORKERS)*kittens::WARP_THREADS, 1) +void fwd_attend_ker(const __grid_constant__ fwd_globals g) { + extern __shared__ int __shm[]; + tma_swizzle_allocator al((int*)&__shm[0]); + int warpid = kittens::warpid(), warpgroupid = warpid/kittens::WARPGROUP_WARPS; + + using K = fwd_attend_ker_tile_dims; + + using q_tile = st_bf; + using k_tile = st_bf; + using v_tile = st_bf; + using l_col_vec = col_vec>; + using o_tile = st_bf; + + q_tile (&q_smem)[CONSUMER_WARPGROUPS] = al.allocate(); + k_tile (&k_smem)[K::stages] = al.allocate(); + v_tile (&v_smem)[K::stages] = al.allocate(); + l_col_vec (&l_smem)[CONSUMER_WARPGROUPS] = al.allocate(); + auto (*o_smem) = reinterpret_cast(q_smem); + int img_kv_blocks; + int kv_blocks = g.N / (K::kv_height); + if constexpr (text_kv) { + img_kv_blocks = kv_blocks - 3; + } else { + img_kv_blocks = kv_blocks; + } + int kv_head_idx = blockIdx.y / g.hr; + int seq_idx; + if constexpr (text_q) { + seq_idx = CT * CH * CW * 6.0 + blockIdx.x * CONSUMER_WARPGROUPS; + } else { + seq_idx = blockIdx.x * CONSUMER_WARPGROUPS; + } + __shared__ kittens::semaphore qsmem_semaphore, k_smem_arrived[K::stages], v_smem_arrived[K::stages], compute_done[K::stages]; + if (threadIdx.x == 0) { + init_semaphore(qsmem_semaphore, 0, 1); + for(int j = 0; j < K::stages; j++) { + init_semaphore(k_smem_arrived[j], 0, 1); + init_semaphore(v_smem_arrived[j], 0, 1); + init_semaphore(compute_done[j], CONSUMER_WARPGROUPS, 0); + } + + tma::expect_bytes(qsmem_semaphore, sizeof(q_smem)); + + for (int wg = 0; wg < CONSUMER_WARPGROUPS; wg++) { + coord q_tile_idx = {blockIdx.z, blockIdx.y, (seq_idx) + wg, 0}; + tma::load_async(q_smem[wg], g.q, q_tile_idx, qsmem_semaphore); + } + + if constexpr (text_q){ + for (int j = 0; j < K::stages - 1; j++) { + coord kv_tile_idx = {blockIdx.z, kv_head_idx, j, 0}; + tma::expect_bytes(k_smem_arrived[j], sizeof(k_tile)); + tma::load_async(k_smem[j], g.k, kv_tile_idx, k_smem_arrived[j]); + tma::expect_bytes(v_smem_arrived[j], sizeof(v_tile)); + tma::load_async(v_smem[j], g.v, kv_tile_idx, v_smem_arrived[j]); + } + } else { + int qt = seq_idx / 6 / (CH * CW); + int qh = (seq_idx / 6) % (CH * CW) / CW; + int qw = (seq_idx / 6) % CW; + qt = clamp_int(qt, DT, CT-DT-1); + qh = clamp_int(qh, DH, CH-DH-1); + qw = clamp_int(qw, DW, CW-DW-1); + int count = 0; + int j = 0; + while (count < K::stages - 1) { + int kt = j / 3 / (CH * CW); + int kh = (j / 3) % (CH * CW) / CW; + int kw = (j / 3) % CW; + bool mask = (abs_int(qt - kt) <= DT) && (abs_int(qh - kh) <= DH) && (abs_int(qw - kw) <= DW); + if (mask){ + coord kv_tile_idx = {blockIdx.z, kv_head_idx, j, 0}; + tma::expect_bytes(k_smem_arrived[count], sizeof(k_tile)); + tma::load_async(k_smem[count], g.k, kv_tile_idx, k_smem_arrived[count]); + tma::expect_bytes(v_smem_arrived[count], sizeof(v_tile)); + tma::load_async(v_smem[count], g.v, kv_tile_idx, v_smem_arrived[count]); + count += 1; + } + j += 1; + } + } + } + __syncthreads(); + + int pipe_idx = K::stages - 1; + + if(warpgroupid == NUM_WARPGROUPS-1) { + warpgroup::decrease_registers<32>(); + + int kv_iters; + if constexpr (is_causal) { + kv_iters = (seq_idx * (K::qo_height/kittens::TILE_ROW_DIM)) - 1 + (CONSUMER_WARPGROUPS * (K::qo_height/kittens::TILE_ROW_DIM)); + kv_iters = ((kv_iters / (K::kv_height/kittens::TILE_ROW_DIM)) == 0) ? (0) : ((kv_iters / (K::kv_height/kittens::TILE_ROW_DIM)) - 1); + } + else { kv_iters = kv_blocks-2;} + + if(warpid == NUM_WORKERS-4) { + if constexpr (text_q){ + for (auto kv_idx = pipe_idx - 1; kv_idx <= kv_iters; kv_idx++) { + coord kv_tile_idx = {blockIdx.z, kv_head_idx, kv_idx + 1, 0}; + tma::expect_bytes(k_smem_arrived[(kv_idx+1)%K::stages], sizeof(k_tile)); + tma::load_async(k_smem[(kv_idx+1)%K::stages], g.k, kv_tile_idx, k_smem_arrived[(kv_idx+1)%K::stages]); + tma::expect_bytes(v_smem_arrived[(kv_idx+1)%K::stages], sizeof(v_tile)); + tma::load_async(v_smem[(kv_idx+1)%K::stages], g.v, kv_tile_idx, v_smem_arrived[(kv_idx+1)%K::stages]); + kittens::wait(compute_done[(kv_idx)%K::stages], (kv_idx/K::stages)%2); + } + } else { + int qt = seq_idx / 6 / (CH * CW); + int qh = (seq_idx / 6) % (CH * CW) / CW; + int qw = (seq_idx / 6) % CW; + qt = clamp_int(qt, DT, CT-DT-1); + qh = clamp_int(qh, DH, CH-DH-1); + qw = clamp_int(qw, DW, CW-DW-1); + int k_t_min = clamp_int(qt-DT, 0, CT-1); + int k_t_max = clamp_int(qt+DT, 0, CT-1); + int k_h_min = clamp_int(qh-DH, 0, CH-1); + int k_h_max = clamp_int(qh+DH, 0, CH-1); + int k_w_min = clamp_int(qw-DW, 0, CW-1); + int k_w_max = clamp_int(qw+DW, 0, CW-1); + int count = 0; + for (int kt = k_t_min; kt <= k_t_max; kt++) { + for (int kh = k_h_min; kh <= k_h_max; kh++) { + for (int kw = k_w_min; kw <= k_w_max; kw++) { + for (int j = 0; j <= 2; j++){ + if (count >= K::stages - 1) { + int index = ((kt * (CH * CW)) + (kh * CW) + kw) * 3 + j; + coord kv_tile_idx = {blockIdx.z, kv_head_idx, index, 0}; + tma::expect_bytes(k_smem_arrived[count%K::stages], sizeof(k_tile)); + tma::load_async(k_smem[count%K::stages], g.k, kv_tile_idx, k_smem_arrived[count%K::stages]); + tma::expect_bytes(v_smem_arrived[count%K::stages], sizeof(v_tile)); + tma::load_async(v_smem[count%K::stages], g.v, kv_tile_idx, v_smem_arrived[count%K::stages]); + kittens::wait(compute_done[(count - 1)%K::stages], ((count - 1)/K::stages)%2); + count += 1; + } else { + count += 1; + } + } + } + } + } + // for text + for (int index = img_kv_blocks; index < kv_blocks; index++) { + coord kv_tile_idx = {blockIdx.z, kv_head_idx, index, 0}; + tma::expect_bytes(k_smem_arrived[count%K::stages], sizeof(k_tile)); + tma::load_async(k_smem[count%K::stages], g.k, kv_tile_idx, k_smem_arrived[count%K::stages]); + tma::expect_bytes(v_smem_arrived[count%K::stages], sizeof(v_tile)); + tma::load_async(v_smem[count%K::stages], g.v, kv_tile_idx, v_smem_arrived[count%K::stages]); + kittens::wait(compute_done[(count - 1)%K::stages], ((count - 1)/K::stages)%2); + count += 1; + } + } + + + } + } + else { + warpgroup::increase_registers<160>(); + + rt_fl<16, K::kv_height> att_block; + rt_bf<16, K::kv_height> att_block_mma; + rt_fl<16, K::tile_width> o_reg; + + col_vec> max_vec, norm_vec, max_vec_last_scaled, max_vec_scaled; + + neg_infty(max_vec); + zero(norm_vec); + zero(o_reg); + + int kv_iters; + if constexpr (is_causal) { + kv_iters = (seq_idx * 4) - 1 + (CONSUMER_WARPGROUPS * 4); + kv_iters = (kv_iters/8); + } + else if constexpr (text_q){ + // the last three kv blocks are for text, we process them separately + kv_iters = img_kv_blocks - 1; + } else { + kv_iters = clamp_int(DT*2+1, 1, CT) * clamp_int(DH*2+1, 1, CH) * clamp_int(DW*2+1, 1, CW) * 3 - 1 ; + } + + kittens::wait(qsmem_semaphore, 0); + for (auto kv_idx = 0; kv_idx <= kv_iters; kv_idx++) { + + kittens::wait(k_smem_arrived[(kv_idx)%K::stages], (kv_idx/K::stages)%2); + warpgroup::mm_ABt(att_block, q_smem[warpgroupid], k_smem[(kv_idx)%K::stages]); + + copy(max_vec_last_scaled, max_vec); + if constexpr (D == 64) { mul(max_vec_last_scaled, max_vec_last_scaled, 1.44269504089f*0.125f); } + else { mul(max_vec_last_scaled, max_vec_last_scaled, 1.44269504089f*0.08838834764f); } + + warpgroup::mma_async_wait(); + + row_max(max_vec, att_block, max_vec); + + if constexpr (D == 64) { + mul(att_block, att_block, 1.44269504089f*0.125f); + mul(max_vec_scaled, max_vec, 1.44269504089f*0.125f); + } + else { + mul(att_block, att_block, 1.44269504089f*0.08838834764f); + mul(max_vec_scaled, max_vec, 1.44269504089f*0.08838834764f); + } + + sub_row(att_block, att_block, max_vec_scaled); + exp2(att_block, att_block); + sub(max_vec_last_scaled, max_vec_last_scaled, max_vec_scaled); + exp2(max_vec_last_scaled, max_vec_last_scaled); + mul(norm_vec, norm_vec, max_vec_last_scaled); + row_sum(norm_vec, att_block, norm_vec); + add(att_block, att_block, 0.f); + copy(att_block_mma, att_block); + mul_row(o_reg, o_reg, max_vec_last_scaled); + + kittens::wait(v_smem_arrived[(kv_idx)%K::stages], (kv_idx/K::stages)%2); + + warpgroup::mma_AB(o_reg, att_block_mma, v_smem[(kv_idx)%K::stages]); + warpgroup::mma_async_wait(); + + if(warpgroup::laneid() == 0) arrive(compute_done[(kv_idx)%K::stages], 1); + } + // the last three kv blocks are for text, we process them separately + if constexpr(text_kv) { + for (auto kv_idx = kv_iters + 1; kv_idx <= kv_iters + 3; kv_idx++) { + + kittens::wait(k_smem_arrived[(kv_idx)%K::stages], (kv_idx/K::stages)%2); + warpgroup::mm_ABt(att_block, q_smem[warpgroupid], k_smem[(kv_idx)%K::stages]); + + copy(max_vec_last_scaled, max_vec); + if constexpr (D == 64) { mul(max_vec_last_scaled, max_vec_last_scaled, 1.44269504089f*0.125f); } + else { mul(max_vec_last_scaled, max_vec_last_scaled, 1.44269504089f*0.08838834764f); } + + warpgroup::mma_async_wait(); + // apply non-pad mask + int offset = g.text_L - (kv_idx - (kv_iters + 1)) * K::kv_height; + // printf("k_idx_start: %d, k_idx_end: %d, text_end: %d, offset: %d\n", k_idx_start, k_idx_end, text_end, offset); + right_fill(att_block, att_block, offset, base_types::constants::neg_infty()); + + + row_max(max_vec, att_block, max_vec); + + if constexpr (D == 64) { + mul(att_block, att_block, 1.44269504089f*0.125f); + mul(max_vec_scaled, max_vec, 1.44269504089f*0.125f); + } + else { + mul(att_block, att_block, 1.44269504089f*0.08838834764f); + mul(max_vec_scaled, max_vec, 1.44269504089f*0.08838834764f); + } + + sub_row(att_block, att_block, max_vec_scaled); + exp2(att_block, att_block); + sub(max_vec_last_scaled, max_vec_last_scaled, max_vec_scaled); + exp2(max_vec_last_scaled, max_vec_last_scaled); + mul(norm_vec, norm_vec, max_vec_last_scaled); + row_sum(norm_vec, att_block, norm_vec); + add(att_block, att_block, 0.f); + copy(att_block_mma, att_block); + mul_row(o_reg, o_reg, max_vec_last_scaled); + + kittens::wait(v_smem_arrived[(kv_idx)%K::stages], (kv_idx/K::stages)%2); + + warpgroup::mma_AB(o_reg, att_block_mma, v_smem[(kv_idx)%K::stages]); + warpgroup::mma_async_wait(); + + if(warpgroup::laneid() == 0) arrive(compute_done[(kv_idx)%K::stages], 1); + } + } + + div_row(o_reg, o_reg, norm_vec); + warpgroup::store(o_smem[warpgroupid], o_reg); + warpgroup::sync(warpgroupid+4); + + if (warpid % 4 == 0) { + coord o_tile_idx = {blockIdx.z, blockIdx.y, (seq_idx) + warpgroupid, 0}; + tma::store_async(g.o, o_smem[warpgroupid], o_tile_idx); + } + + mul(max_vec_scaled, max_vec_scaled, 0.69314718056f); + log(norm_vec, norm_vec); + add(norm_vec, norm_vec, max_vec_scaled); + + if constexpr (D == 64) { mul(norm_vec, norm_vec, -8.0f); } + else { mul(norm_vec, norm_vec, -11.313708499f); } + + warpgroup::store(l_smem[warpgroupid], norm_vec); + warpgroup::sync(warpgroupid+4); + + if (warpid % 4 == 0) { + coord tile_idx = {blockIdx.z, blockIdx.y, 0, (seq_idx) + warpgroupid}; + tma::store_async(g.l, l_smem[warpgroupid], tile_idx); + } + tma::store_async_wait(); + } +} + + + +#include "pyutils/torch_helpers.cuh" +#include +#include + +torch::Tensor +sta_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o, int kernel_t_size, int kernel_h_size, int kernel_w_size, int text_length, bool process_text, bool has_text, int kernel_aspect_ratio_flag) +{ + CHECK_INPUT(q); + CHECK_INPUT(k); + CHECK_INPUT(v); + + auto batch = q.size(0); + auto seq_len = q.size(2); + auto head_dim = q.size(3); + auto qo_heads = q.size(1); + auto kv_heads = k.size(1); + + // check to see that these dimensions match for all inputs + TORCH_CHECK(q.size(0) == batch, "Q batch dimension - idx 0 - must match for all inputs"); + TORCH_CHECK(k.size(0) == batch, "K batch dimension - idx 0 - must match for all inputs"); + TORCH_CHECK(v.size(0) == batch, "V batch dimension - idx 0 - must match for all inputs"); + + TORCH_CHECK(q.size(2) == seq_len, "Q sequence length dimension - idx 2 - must match for all inputs"); + TORCH_CHECK(k.size(2) == seq_len, "K sequence length dimension - idx 2 - must match for all inputs"); + TORCH_CHECK(v.size(2) == seq_len, "V sequence length dimension - idx 2 - must match for all inputs"); + + TORCH_CHECK(q.size(3) == head_dim, "Q head dimension - idx 3 - must match for all non-vector inputs"); + TORCH_CHECK(k.size(3) == head_dim, "K head dimension - idx 3 - must match for all non-vector inputs"); + TORCH_CHECK(v.size(3) == head_dim, "V head dimension - idx 3 - must match for all non-vector inputs"); + + TORCH_CHECK(qo_heads >= kv_heads, "QO heads must be greater than or equal to KV heads"); + TORCH_CHECK(qo_heads % kv_heads == 0, "QO heads must be divisible by KV heads"); + TORCH_CHECK(q.size(1) == qo_heads, "QO head dimension - idx 1 - must match for all inputs"); + TORCH_CHECK(k.size(1) == kv_heads, "KV head dimension - idx 1 - must match for all inputs"); + TORCH_CHECK(v.size(1) == kv_heads, "KV head dimension - idx 1 - must match for all inputs"); + + auto hr = qo_heads / kv_heads; + + c10::BFloat16* q_ptr = q.data_ptr(); + c10::BFloat16* k_ptr = k.data_ptr(); + c10::BFloat16* v_ptr = v.data_ptr(); + + bf16* d_q = reinterpret_cast(q_ptr); + bf16* d_k = reinterpret_cast(k_ptr); + bf16* d_v = reinterpret_cast(v_ptr); + + + + torch::Tensor l_vec = torch::empty({static_cast(batch), + static_cast(qo_heads), + static_cast(seq_len), + static_cast(1)}, + torch::TensorOptions().dtype(torch::kFloat).device(q.device()).memory_format(at::MemoryFormat::Contiguous)); + + + bf16* o_ptr = reinterpret_cast(o.data_ptr()); + bf16* d_o = reinterpret_cast(o_ptr); + + float* l_ptr = reinterpret_cast(l_vec.data_ptr()); + float* d_l = reinterpret_cast(l_ptr); + + //cudadevicesynchronize(); + const c10::cuda::OptionalCUDAGuard device_guard(q.device()); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); + + + if (head_dim == 128) { + using q_tile = st_bf::qo_height, fwd_attend_ker_tile_dims<128>::tile_width>; + using k_tile = st_bf::kv_height, fwd_attend_ker_tile_dims<128>::tile_width>; + using v_tile = st_bf::kv_height, fwd_attend_ker_tile_dims<128>::tile_width>; + using l_col_vec = col_vec::qo_height, fwd_attend_ker_tile_dims<128>::tile_width>>; + using o_tile = st_bf::qo_height, fwd_attend_ker_tile_dims<128>::tile_width>; + + using q_global = gl; + using k_global = gl; + using v_global = gl; + using l_global = gl; + using o_global = gl; + + using globals = fwd_globals<128>; + + q_global qg_arg{d_q, static_cast(batch), static_cast(qo_heads), static_cast(seq_len), 128U}; + k_global kg_arg{d_k, static_cast(batch), static_cast(kv_heads), static_cast(seq_len), 128U}; + v_global vg_arg{d_v, static_cast(batch), static_cast(kv_heads), static_cast(seq_len), 128U}; + l_global lg_arg{d_l, static_cast(batch), static_cast(qo_heads), 1U, static_cast(seq_len)}; + o_global og_arg{d_o, static_cast(batch), static_cast(qo_heads), static_cast(seq_len), 128U}; + + globals g{qg_arg, kg_arg, vg_arg, lg_arg, og_arg, static_cast(seq_len), static_cast(text_length), static_cast(hr)}; + + // Shared memory size for the kernel. + // We use the maximum available shared memory (kittens::MAX_SHARED_MEMORY) + // which is approximately 227KB on H100, necessary for the high-performance + // TMA-based attention tiles with multiple stages. + constexpr int mem_size = kittens::MAX_SHARED_MEMORY; + int threads = NUM_WORKERS * kittens::WARP_THREADS; + if (has_text) { + // TORCH_CHECK(seq_len % (CONSUMER_WARPGROUPS*kittens::TILE_DIM*4) == 0, "sequence length must be divisible by 192"); + dim3 grid_image(seq_len/(CONSUMER_WARPGROUPS*kittens::TILE_ROW_DIM*4)-2, qo_heads, batch); + dim3 grid_text(2, qo_heads, batch); + if (!process_text) { +#define LAUNCH_IMAGE_KER(DT_VAL, DH_VAL, DW_VAL) \ + cudaFuncSetAttribute( \ + fwd_attend_ker<128, false, false, true, DT_VAL, DH_VAL, DW_VAL, 5, 6, 10>, \ + cudaFuncAttributeMaxDynamicSharedMemorySize, \ + mem_size \ + ); \ + fwd_attend_ker<128, false, false, true, DT_VAL, DH_VAL, DW_VAL, 5, 6, 10><<>>(g); + + if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 3) { LAUNCH_IMAGE_KER(1, 1, 1); } + else if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 5) { LAUNCH_IMAGE_KER(1, 1, 2); } + else if (kernel_t_size == 5 && kernel_h_size == 3 && kernel_w_size == 3) { LAUNCH_IMAGE_KER(2, 1, 1); } + else if (kernel_t_size == 3 && kernel_h_size == 5 && kernel_w_size == 5) { LAUNCH_IMAGE_KER(1, 2, 2); } + else if (kernel_t_size == 5 && kernel_h_size == 6 && kernel_w_size == 1) { LAUNCH_IMAGE_KER(2, 3, 0); } + else if (kernel_t_size == 5 && kernel_h_size == 3 && kernel_w_size == 5) { LAUNCH_IMAGE_KER(2, 1, 2); } + else if (kernel_t_size == 5 && kernel_h_size == 5 && kernel_w_size == 5) { LAUNCH_IMAGE_KER(2, 2, 2); } + else if (kernel_t_size == 5 && kernel_h_size == 5 && kernel_w_size == 7) { LAUNCH_IMAGE_KER(2, 2, 3); } + else if (kernel_t_size == 5 && kernel_h_size == 6 && kernel_w_size == 10){ LAUNCH_IMAGE_KER(2, 3, 5); } + else if (kernel_t_size == 3 && kernel_h_size == 6 && kernel_w_size == 10){ LAUNCH_IMAGE_KER(1, 3, 5); } + else if (kernel_t_size == 5 && kernel_h_size == 1 && kernel_w_size == 1) { LAUNCH_IMAGE_KER(2, 0, 0); } + else if (kernel_t_size == 1 && kernel_h_size == 6 && kernel_w_size == 10){ LAUNCH_IMAGE_KER(0, 3, 5); } + else if (kernel_t_size == 5 && kernel_h_size == 1 && kernel_w_size == 10){ LAUNCH_IMAGE_KER(2, 0, 5); } + else { + TORCH_CHECK(false, "Invalid kernel size: ", kernel_t_size, "x", kernel_h_size, "x", kernel_w_size); + } +#undef LAUNCH_IMAGE_KER + } else { + cudaFuncSetAttribute( + fwd_attend_ker<128, false, true, true, 1, 1, 1, 5, 6, 10>, + cudaFuncAttributeMaxDynamicSharedMemorySize, + mem_size + ); + fwd_attend_ker<128, false, true, true, 1, 1, 1, 5, 6, 10><<>>(g); + } + + } else { + dim3 grid_image(seq_len/(CONSUMER_WARPGROUPS*kittens::TILE_ROW_DIM*4), qo_heads, batch); + if (kernel_aspect_ratio_flag == 2){ +#define LAUNCH_IMAGE_KER(DT_VAL, DH_VAL, DW_VAL) \ + cudaFuncSetAttribute( \ + fwd_attend_ker<128, false, false, false, DT_VAL, DH_VAL, DW_VAL, 6, 6, 6>, \ + cudaFuncAttributeMaxDynamicSharedMemorySize, \ + mem_size \ + ); \ + fwd_attend_ker<128, false, false, false, DT_VAL, DH_VAL, DW_VAL, 6, 6, 6><<>>(g); + + if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 3) { LAUNCH_IMAGE_KER(1, 1, 1); } + else if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 6) { LAUNCH_IMAGE_KER(1, 1, 3); } + else if (kernel_t_size == 6 && kernel_h_size == 3 && kernel_w_size == 3) { LAUNCH_IMAGE_KER(3, 1, 1); } + else if (kernel_t_size == 3 && kernel_h_size == 6 && kernel_w_size == 6) { LAUNCH_IMAGE_KER(1, 3, 3); } + else if (kernel_t_size == 3 && kernel_h_size == 6 && kernel_w_size == 3) { LAUNCH_IMAGE_KER(1, 3, 1); } + else if (kernel_t_size == 6 && kernel_h_size == 3 && kernel_w_size == 6) { LAUNCH_IMAGE_KER(3, 1, 3); } + else if (kernel_t_size == 6 && kernel_h_size == 6 && kernel_w_size == 6) { LAUNCH_IMAGE_KER(3, 3, 3); } + else if (kernel_t_size == 6 && kernel_h_size == 1 && kernel_w_size == 1) { LAUNCH_IMAGE_KER(3, 0, 0); } + else if (kernel_t_size == 6 && kernel_h_size == 1 && kernel_w_size == 6) { LAUNCH_IMAGE_KER(3, 0, 3); } + else if (kernel_t_size == 6 && kernel_h_size == 6 && kernel_w_size == 1) { LAUNCH_IMAGE_KER(3, 3, 0); } + else if (kernel_t_size == 1 && kernel_h_size == 6 && kernel_w_size == 6) { LAUNCH_IMAGE_KER(0, 3, 3); } + else if (kernel_t_size == 1 && kernel_h_size == 1 && kernel_w_size == 6) { LAUNCH_IMAGE_KER(0, 0, 3); } + else if (kernel_t_size == 1 && kernel_h_size == 6 && kernel_w_size == 1) { LAUNCH_IMAGE_KER(0, 3, 0); } + else { + TORCH_CHECK(false, "Invalid kernel size: ", kernel_t_size, "x", kernel_h_size, "x", kernel_w_size); + } +#undef LAUNCH_IMAGE_KER + } + else if (kernel_aspect_ratio_flag == 3) { +#define LAUNCH_IMAGE_KER(DT_VAL, DH_VAL, DW_VAL) \ + cudaFuncSetAttribute( \ + fwd_attend_ker<128, false, false, false, DT_VAL, DH_VAL, DW_VAL, 3, 6, 10>, \ + cudaFuncAttributeMaxDynamicSharedMemorySize, \ + mem_size \ + ); \ + fwd_attend_ker<128, false, false, false, DT_VAL, DH_VAL, DW_VAL, 3, 6, 10><<>>(g); + + if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 3) { LAUNCH_IMAGE_KER(1, 1, 1); } + else if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 5) { LAUNCH_IMAGE_KER(1, 1, 2); } + else if (kernel_t_size == 3 && kernel_h_size == 5 && kernel_w_size == 5) { LAUNCH_IMAGE_KER(1, 2, 2); } + else if (kernel_t_size == 3 && kernel_h_size == 6 && kernel_w_size == 1) { LAUNCH_IMAGE_KER(1, 3, 0); } + else if (kernel_t_size == 3 && kernel_h_size == 5 && kernel_w_size == 7) { LAUNCH_IMAGE_KER(1, 2, 3); } + else if (kernel_t_size == 3 && kernel_h_size == 5 && kernel_w_size == 9) { LAUNCH_IMAGE_KER(1, 2, 4); } + else if (kernel_t_size == 3 && kernel_h_size == 6 && kernel_w_size == 10){ LAUNCH_IMAGE_KER(1, 3, 5); } + else if (kernel_t_size == 3 && kernel_h_size == 6 && kernel_w_size == 3) { LAUNCH_IMAGE_KER(1, 3, 1); } + else if (kernel_t_size == 3 && kernel_h_size == 1 && kernel_w_size == 1) { LAUNCH_IMAGE_KER(1, 0, 0); } + else if (kernel_t_size == 1 && kernel_h_size == 6 && kernel_w_size == 10){ LAUNCH_IMAGE_KER(0, 3, 5); } + else if (kernel_t_size == 1 && kernel_h_size == 5 && kernel_w_size == 10){ LAUNCH_IMAGE_KER(0, 2, 5); } + else if (kernel_t_size == 1 && kernel_h_size == 6 && kernel_w_size == 7) { LAUNCH_IMAGE_KER(0, 3, 3); } + else if (kernel_t_size == 1 && kernel_h_size == 5 && kernel_w_size == 7) { LAUNCH_IMAGE_KER(0, 2, 3); } + else if (kernel_t_size == 1 && kernel_h_size == 5 && kernel_w_size == 9) { LAUNCH_IMAGE_KER(0, 2, 4); } + else if (kernel_t_size == 3 && kernel_h_size == 1 && kernel_w_size == 10){ LAUNCH_IMAGE_KER(1, 0, 5); } + else if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 10){ LAUNCH_IMAGE_KER(1, 1, 5); } + else if (kernel_t_size == 1 && kernel_h_size == 3 && kernel_w_size == 10){ LAUNCH_IMAGE_KER(0, 1, 5); } + else if (kernel_t_size == 1 && kernel_h_size == 6 && kernel_w_size == 5) { LAUNCH_IMAGE_KER(0, 3, 2); } + else { + TORCH_CHECK(false, "Invalid kernel size: ", kernel_t_size, "x", kernel_h_size, "x", kernel_w_size); + } +#undef LAUNCH_IMAGE_KER + } + + else { + TORCH_CHECK(false, "Unsupported kernel_aspect_ratio_flag: ", kernel_aspect_ratio_flag); + } + + } + CHECK_CUDA_ERROR(cudaGetLastError()); + // cudaStreamSynchronize(stream); + } + + return o; + //cudadevicesynchronize(); +} + diff --git a/csrc/fastvideo_kernel/csrc/vsa.cpp b/csrc/fastvideo_kernel/csrc/vsa.cpp new file mode 100644 index 000000000..8661103a0 --- /dev/null +++ b/csrc/fastvideo_kernel/csrc/vsa.cpp @@ -0,0 +1,27 @@ +#include +#include + +#include +#include +#include + +#include + +#ifdef TK_COMPILE_BLOCK_SPARSE +extern std::vector block_sparse_attention_forward( + torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor q2k_block_sparse_index, torch::Tensor q2k_block_sparse_num, torch::Tensor block_size +); +extern std::vector block_sparse_attention_backward( + torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o, torch::Tensor l_vec, torch::Tensor og, torch::Tensor k2q_block_sparse_index, torch::Tensor k2q_block_sparse_num, torch::Tensor block_size +); +#endif + + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.doc() = "Video Sparse Attention Kernels"; // optional module docstring + +#ifdef TK_COMPILE_BLOCK_SPARSE + m.def("block_sparse_fwd", torch::wrap_pybind_function(block_sparse_attention_forward), "block sparse attention"); + m.def("block_sparse_bwd", torch::wrap_pybind_function(block_sparse_attention_backward), "block sparse attention backward"); +#endif +} diff --git a/csrc/fastvideo_kernel/pyproject.toml b/csrc/fastvideo_kernel/pyproject.toml new file mode 100644 index 000000000..2d2f72748 --- /dev/null +++ b/csrc/fastvideo_kernel/pyproject.toml @@ -0,0 +1,27 @@ +[build-system] +requires = ["setuptools>=61.0", "torch>=2.5.0", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "fastvideo-kernel" +version = "0.1.0" +description = "CUDA kernels for FastVideo" +readme = "README.md" +requires-python = ">=3.10" +license = {text = "Apache-2.0"} +authors = [{name = "Hao AI Lab"}] +classifiers = [ + "Programming Language :: Python :: 3", + "Environment :: GPU :: NVIDIA CUDA :: 12", + "License :: OSI Approved :: Apache Software License", +] +dependencies = [ + "torch>=2.5.0", + "triton>=2.0.0" +] + +[project.urls] +Repository = "https://github.com/hao-ai-lab/FastVideo" + +[tool.setuptools.packages.find] +where = ["src"] diff --git a/csrc/fastvideo_kernel/setup.py b/csrc/fastvideo_kernel/setup.py new file mode 100644 index 000000000..455306704 --- /dev/null +++ b/csrc/fastvideo_kernel/setup.py @@ -0,0 +1,132 @@ +import os +import subprocess +import sys +from pathlib import Path +from setuptools import find_packages, setup +from torch.utils.cpp_extension import BuildExtension, CUDAExtension + +ROOT = Path(__file__).parent.absolute() +CSRC_DIR = ROOT / "csrc" + +# Path to ThunderKittens (TK) +def get_tk_dir(): + tk_env = os.getenv("THUNDERKITTENS_ROOT") + if tk_env: + return tk_env + + # Check common locations + possible_paths = [ + ROOT / "tk", + ROOT / "csrc" / "tk", + ROOT.parent / "attn" / "sliding_tile_attn" / "tk", + ROOT.parent / "attn" / "video_sparse_attn" / "tk", + ] + for p in possible_paths: + if (p / "include" / "kittens.cuh").exists(): + return str(p) + + # Default fallback + return str(ROOT.parent / "attn" / "sliding_tile_attn" / "tk") + +TK_DIR = get_tk_dir() + +def get_cuda_flags(tk_root: str) -> list: + python_include = subprocess.check_output( + ["python", "-c", "import sysconfig; print(sysconfig.get_path('include'))"] + ).decode().strip() + + torch_includes = subprocess.check_output([ + "python", "-c", + "import torch; from torch.utils.cpp_extension import include_paths; " + "print(' '.join(['-I' + p for p in include_paths()]))" + ]).decode().strip().split() + + return [ + "-DNDEBUG", + "-Xcompiler=-Wno-psabi", + "-Xcompiler=-fno-strict-aliasing", + "--expt-extended-lambda", + "--expt-relaxed-constexpr", + "-forward-unknown-to-host-compiler", + "--use_fast_math", + "-std=c++20", + "-O3", + "-Xnvlink=--verbose", + "-Xptxas=--verbose", + "-Xptxas=--warn-on-spills", + f"-I{tk_root}/include", + f"-I{tk_root}/prototype", + f"-I{python_include}", + "-DTORCH_COMPILE", + "-DKITTENS_HOPPER", + "-arch=sm_90a", + ] + torch_includes + +def get_extensions(): + if not torch.cuda.is_available(): + return [] + + extensions = [] + cpp_flags = ["-std=c++20", "-O3"] + + # Check if TK is available + if not os.path.exists(os.path.join(TK_DIR, "include", "kittens.cuh")): + print(f"Warning: ThunderKittens not found at {TK_DIR}. CUDA kernels will not be built.") + return [] + + cuda_flags = get_cuda_flags(TK_DIR) + + # STA Extension + extensions.append(CUDAExtension( + "fastvideo_kernel._C.st_attn", + sources=[ + "csrc/st_attn.cpp", + "csrc/st_attn_h100.cu", + ], + extra_compile_args={ + "cxx": cpp_flags + ["-DTK_COMPILE_ST_ATTN"], + "nvcc": cuda_flags + ["-DTK_COMPILE_ST_ATTN"] + }, + libraries=["cuda"], + )) + + # VSA Extension + extensions.append(CUDAExtension( + "fastvideo_kernel._C.vsa", + sources=[ + "csrc/vsa.cpp", + "csrc/block_sparse_h100.cu", + ], + extra_compile_args={ + "cxx": cpp_flags + ["-DTK_COMPILE_BLOCK_SPARSE"], + "nvcc": cuda_flags + ["-DTK_COMPILE_BLOCK_SPARSE"] + }, + libraries=["cuda"], + )) + + return extensions + +ext_modules = [] +if not any(arg in sys.argv for arg in ["clean", "egg_info", "--version"]): + try: + import torch + ext_modules = get_extensions() + except Exception as e: + print(f"Warning: Failed to configure CUDA extensions: {e}") + +setup( + name="fastvideo-kernel", + version="0.1.0", + description="Unified CUDA kernels for FastVideo", + long_description=open("README.md").read(), + long_description_content_type="text/markdown", + license="Apache-2.0", + author="Hao AI Lab", + url="https://github.com/hao-ai-lab/FastVideo", + package_dir={"": "src"}, + packages=find_packages(where="src"), + ext_modules=ext_modules, + cmdclass={"build_ext": BuildExtension} if ext_modules else {}, + python_requires=">=3.10", + install_requires=["torch>=2.5.0", "triton>=2.0.0"], +) diff --git a/csrc/fastvideo_kernel/src/fastvideo_kernel/_C/__init__.py b/csrc/fastvideo_kernel/src/fastvideo_kernel/_C/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/csrc/fastvideo_kernel/src/fastvideo_kernel/__init__.py b/csrc/fastvideo_kernel/src/fastvideo_kernel/__init__.py new file mode 100644 index 000000000..cd4a74f1d --- /dev/null +++ b/csrc/fastvideo_kernel/src/fastvideo_kernel/__init__.py @@ -0,0 +1,21 @@ +__version__ = "0.1.0" + +from fastvideo_kernel.ops import ( + sliding_tile_attention, + video_sparse_attn, +) + +from fastvideo_kernel.vmoba import ( + moba_attn_varlen, + process_moba_input, + process_moba_output, +) + +__all__ = [ + "sliding_tile_attention", + "video_sparse_attn", + "moba_attn_varlen", + "process_moba_input", + "process_moba_output", + "__version__", +] diff --git a/csrc/fastvideo_kernel/src/fastvideo_kernel/ops.py b/csrc/fastvideo_kernel/src/fastvideo_kernel/ops.py new file mode 100644 index 000000000..b5feeacdc --- /dev/null +++ b/csrc/fastvideo_kernel/src/fastvideo_kernel/ops.py @@ -0,0 +1,103 @@ +import math +import torch +from .triton_kernels.block_sparse_attn_triton import triton_block_sparse_attn_forward +from .triton_kernels.index import map_to_index + +try: + from fastvideo_kernel._C.st_attn import sta_fwd +except ImportError: + sta_fwd = None + +try: + from fastvideo_kernel._C.vsa import block_sparse_fwd, block_sparse_bwd +except ImportError: + block_sparse_fwd = None + block_sparse_bwd = None + + +def sliding_tile_attention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + window_size: list, + text_length: int, + has_text: bool = True, + seq_shape: str = "30x48x80", +) -> torch.Tensor: + if sta_fwd is None: + raise RuntimeError("STA kernel not compiled. Requires H100 and ThunderKittens at build time.") + + seq_length = q.shape[2] + shape_map = {"30x48x80": 1, "36x48x48": 2, "18x48x80": 3} + + if has_text: + target_size = math.ceil(seq_length / 384) * 384 + pad_size = target_size - seq_length + if pad_size > 0: + q = torch.cat([q, q[:, :, -pad_size:]], dim=2) + k = torch.cat([k, k[:, :, -pad_size:]], dim=2) + v = torch.cat([v, v[:, :, -pad_size:]], dim=2) + + output = torch.empty_like(q) + flag = shape_map[seq_shape] + + for head_idx, (t, h, w) in enumerate(window_size): + sta_fwd( + q[:, head_idx:head_idx+1], + k[:, head_idx:head_idx+1], + v[:, head_idx:head_idx+1], + output[:, head_idx:head_idx+1], + t, h, w, text_length, False, has_text, flag + ) + + if has_text: + sta_fwd(q, k, v, output, 3, 3, 3, text_length, True, True, flag) + + return output[:, :, :seq_length] + + +def video_sparse_attn( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + variable_block_sizes: torch.Tensor, + topk: int, + block_size: int | tuple = 64, + compress_attn_weight: torch.Tensor = None, +) -> torch.Tensor: + if isinstance(block_size, int): + block_size = (block_size, block_size, block_size) + + block_elements = block_size[0] * block_size[1] * block_size[2] + batch, heads, seq_len, dim = q.shape + + # Compression branch + q_c = q.view(batch, heads, seq_len // block_elements, block_elements, dim) + k_c = k.view(batch, heads, seq_len // block_elements, block_elements, dim) + v_c = v.view(batch, heads, seq_len // block_elements, block_elements, dim) + + q_c = (q_c.float().sum(dim=3) / variable_block_sizes.view(1, 1, -1, 1)).to(q.dtype) + k_c = (k_c.float().sum(dim=3) / variable_block_sizes.view(1, 1, -1, 1)).to(k.dtype) + v_c = (v_c.float().sum(dim=3) / variable_block_sizes.view(1, 1, -1, 1)).to(v.dtype) + + scores = torch.matmul(q_c, k_c.transpose(-2, -1)) / (dim ** 0.5) + attn = torch.softmax(scores, dim=-1) + out_c = torch.matmul(attn, v_c) + + out_c = out_c.view(batch, heads, seq_len // block_elements, 1, dim) + out_c = out_c.repeat(1, 1, 1, block_elements, 1).view(batch, heads, seq_len, dim) + + # Sparse branch + topk_idx = torch.topk(scores, topk, dim=-1).indices + mask = torch.zeros_like(scores, dtype=torch.bool).scatter_(-1, topk_idx, True) + + if block_sparse_fwd is not None: + idx, num = map_to_index(mask) + out_s, _ = block_sparse_fwd(q, k, v, idx, num, variable_block_sizes.int()) + else: + idx, num = map_to_index(mask) + out_s, _ = triton_block_sparse_attn_forward(q, k, v, idx, num, variable_block_sizes) + + if compress_attn_weight is not None: + return out_c * compress_attn_weight + out_s + return out_c + out_s diff --git a/csrc/fastvideo_kernel/src/fastvideo_kernel/triton_kernels/block_sparse_attn_triton.py b/csrc/fastvideo_kernel/src/fastvideo_kernel/triton_kernels/block_sparse_attn_triton.py new file mode 100644 index 000000000..373e2f2ae --- /dev/null +++ b/csrc/fastvideo_kernel/src/fastvideo_kernel/triton_kernels/block_sparse_attn_triton.py @@ -0,0 +1,449 @@ +""" +Fused Attention +=============== + +This is a Triton implementation of the Flash Attention v2 algorithm from Tri Dao +(https://tridao.me/publications/flash2/flash2.pdf) + +Credits: OpenAI kernel team +""" + + +import torch +import triton +import triton.language as tl + +# ──────────────────────────── SPARSE ADDITION BEGIN ─────────────────────────── +import math # small utility needed by the sparse wrapper +# ──────────────────────────── SPARSE ADDITION END ───────────────────────────── + + +# We don't run auto-tuning every time to keep the tutorial fast. Keeping +# the code below and commenting out the equivalent parameters is convenient for +# re-tuning. +configs = [ + triton.Config({'BLOCK_M': BM, 'BLOCK_N': BN}, num_stages=s, num_warps=w) \ + for BM in [64]\ + for BN in [64]\ + for s in [3, 4, 7]\ + for w in [4, 8]\ +] + +# ──────────────────────────── SPARSE ADDITION BEGIN ─────────────────────────── +@triton.autotune(configs, key=["N_CTX", "HEAD_DIM"]) +@triton.jit +def _attn_fwd_sparse(Q, K, V, sm_scale, # + q2k_index, q2k_num, max_kv_blks, # + variable_block_sizes, + M, Out, # + stride_qz, stride_qh, stride_qm, stride_qk, + stride_kz, stride_kh, stride_kn, stride_kk, + stride_vz, stride_vh, stride_vk, stride_vn, + stride_oz, stride_oh, stride_om, stride_on, + Z, H, N_CTX, # + HEAD_DIM: tl.constexpr, # + BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, + STAGE: tl.constexpr): + """ + 64×64 **block-sparse** forward kernel. Back-prop kernels remain dense + (32×64 and 64×32) – memory footprint unchanged. + """ + + # ----- program-id mapping ----- + q_blk = tl.program_id(0) # Q-tile index + off_hz = tl.program_id(1) # fused (batch, head) + b = off_hz // H + h = off_hz % H + q_tiles = N_CTX // BLOCK_M + meta_base = ((b * H + h) * q_tiles + q_blk) + + kv_blocks = tl.load(q2k_num + meta_base) # int32 + kv_ptr = q2k_index + meta_base * max_kv_blks # ptr to list + + # ----- base pointers ----- + qvk_off = (b.to(tl.int64) * stride_qz + + h.to(tl.int64) * stride_qh) + + Q_ptr = tl.make_block_ptr( + base=Q + qvk_off, shape=(N_CTX, HEAD_DIM), + strides=(stride_qm, stride_qk), + offsets=(q_blk * BLOCK_M, 0), + block_shape=(BLOCK_M, HEAD_DIM), order=(1, 0)) + + K_base = tl.make_block_ptr( + base=K + qvk_off, shape=(HEAD_DIM, N_CTX), + strides=(stride_kk, stride_kn), + offsets=(0, 0), + block_shape=(HEAD_DIM, BLOCK_N), order=(0, 1)) + + v_order: tl.constexpr = (0, 1) if V.dtype.element_ty == tl.float8e5 else (1, 0) + V_base = tl.make_block_ptr( + base=V + qvk_off, shape=(N_CTX, HEAD_DIM), + strides=(stride_vk, stride_vn), + offsets=(0, 0), + block_shape=(BLOCK_N, HEAD_DIM), order=v_order) + + O_ptr = tl.make_block_ptr( + base=Out + qvk_off, shape=(N_CTX, HEAD_DIM), + strides=(stride_om, stride_on), + offsets=(q_blk * BLOCK_M, 0), + block_shape=(BLOCK_M, HEAD_DIM), order=(1, 0)) + + # ----- accumulators ----- + offs_m = q_blk * BLOCK_M + tl.arange(0, BLOCK_M) + m_i = tl.full([BLOCK_M], -float("inf"), tl.float32) + l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + 1.0 + acc = tl.zeros([BLOCK_M, HEAD_DIM], dtype=tl.float32) + qk_scale = sm_scale * 1.44269504 # 1/ln2 + q = tl.load(Q_ptr) + + # ----- sparse loop over valid K/V tiles ----- + for i in range(0, kv_blocks): + kv_idx = tl.load(kv_ptr + i).to(tl.int32) + block_size = tl.load(variable_block_sizes + kv_idx) + K_ptr = tl.advance(K_base, (0, kv_idx * BLOCK_N)) + V_ptr = tl.advance(V_base, (kv_idx * BLOCK_N, 0)) + + k = tl.load(K_ptr) + qk = tl.dot(q, k) + # mask out invalid columns + mask = tl.arange(0, BLOCK_N) < block_size + qk = tl.where(mask[None, :], qk, -float("inf")) + + m_ij = tl.maximum(m_i, tl.max(qk, 1) * qk_scale) + p = tl.math.exp2(qk * qk_scale - m_ij[:, None]) + l_ij = tl.sum(p, 1) + + alpha = tl.math.exp2(m_i - m_ij) + l_i = l_i * alpha + l_ij + acc = acc * alpha[:, None] + + v = tl.load(V_ptr) + acc = tl.dot(p.to(tl.bfloat16), v, acc) + m_i = m_ij + + # ----- epilogue ----- + m_i += tl.math.log2(l_i) + acc = acc / l_i[:, None] + tl.store(M + off_hz * N_CTX + offs_m, m_i) + tl.store(O_ptr, acc.to(Out.type.element_ty)) +# ──────────────────────────── SPARSE ADDITION END ───────────────────────────── + + + + +@triton.jit +def _attn_bwd_preprocess(O, DO, # + Delta, # + Z, H, N_CTX, # + BLOCK_M: tl.constexpr, HEAD_DIM: tl.constexpr # + ): + off_m = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M) + off_hz = tl.program_id(1) + off_n = tl.arange(0, HEAD_DIM) + # load + o = tl.load(O + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM + off_n[None, :]) + do = tl.load(DO + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM + off_n[None, :]).to(tl.float32) + delta = tl.sum(o * do, axis=1) + # write-back + tl.store(Delta + off_hz * N_CTX + off_m, delta) + + +# The main inner-loop logic for computing dK and dV. +@triton.jit +def _attn_bwd_dkdv(dk, dv, # + Q, k, v, sm_scale, # + DO, # + M, D, # + k2q_index, k2q_num, max_q_blks, + variable_block_sizes, + # shared by Q/K/V/DO. + stride_tok, stride_d, # + H, N_CTX, BLOCK_M1: tl.constexpr, # + BLOCK_N1: tl.constexpr, # + HEAD_DIM: tl.constexpr, # + # Filled in by the wrapper. + start_n, start_m, num_steps): + offs_m = start_m + tl.arange(0, BLOCK_M1) + offs_n = start_n + tl.arange(0, BLOCK_N1) + offs_k = tl.arange(0, HEAD_DIM) + qT_ptrs = Q + offs_m[None, :] * stride_tok + offs_k[:, None] * stride_d + do_ptrs = DO + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d + # BLOCK_N1 must be a multiple of BLOCK_M1, otherwise the code wouldn't work. + tl.static_assert(BLOCK_N1 % BLOCK_M1 == 0) + step_m = BLOCK_M1 + kv_blk = tl.program_id(0) # Q-tile index + off_hz = tl.program_id(2) # fused (batch, head) + b = off_hz // H + h = off_hz % H + q_tiles = N_CTX // BLOCK_N1 + meta_base = ((b * H + h) * q_tiles + kv_blk) + + q_blocks = tl.load(k2q_num + meta_base) # int32 + q_ptr = k2q_index + meta_base * max_q_blks # ptr to list + block_size = tl.load(variable_block_sizes + kv_blk) + + + + for blk_idx in range(q_blocks*2): + block_sparse_offset = (tl.load(q_ptr + blk_idx//2).to(tl.int32)*2 + blk_idx%2) *step_m + qT = tl.load(qT_ptrs + block_sparse_offset * stride_tok) + # Load m before computing qk to reduce pipeline stall. + offs_m = start_m + block_sparse_offset + tl.arange(0, BLOCK_M1) + m = tl.load(M + offs_m) + qkT = tl.dot(k, qT) + pT = tl.math.exp2(qkT - m[None, :]) + mask = tl.arange(0, BLOCK_N1) < block_size + pT = tl.where(mask[:, None], pT, 0.0) + + do = tl.load(do_ptrs + block_sparse_offset * stride_tok) + # Compute dV. + ppT = pT + ppT = ppT.to(tl.bfloat16) + dv += tl.dot(ppT, do) + # D (= delta) is pre-divided by ds_scale. + Di = tl.load(D + offs_m) + # Compute dP and dS. + dpT = tl.dot(v, tl.trans(do)).to(tl.float32) + dsT = pT * (dpT - Di[None, :]) + dsT = dsT.to(tl.bfloat16) + dk += tl.dot(dsT, tl.trans(qT)) + # Increment pointers. + return dk, dv + + + +# the main inner-loop logic for computing dQ +@triton.jit +def _attn_bwd_dq(dq, q, K, V, # + do, m, D, + # shared by Q/K/V/DO. + q2k_index, q2k_num, max_kv_blks, + variable_block_sizes, + stride_tok, stride_d, # + H, N_CTX, # + BLOCK_M2: tl.constexpr, # + BLOCK_N2: tl.constexpr, # + HEAD_DIM: tl.constexpr, + # Filled in by the wrapper. + start_m, start_n, num_steps): + offs_m = start_m + tl.arange(0, BLOCK_M2) + offs_n = start_n + tl.arange(0, BLOCK_N2) + offs_k = tl.arange(0, HEAD_DIM) + kT_ptrs = K + offs_n[None, :] * stride_tok + offs_k[:, None] * stride_d + vT_ptrs = V + offs_n[None, :] * stride_tok + offs_k[:, None] * stride_d + # D (= delta) is pre-divided by ds_scale. + Di = tl.load(D + offs_m) + # BLOCK_M2 must be a multiple of BLOCK_N2, otherwise the code wouldn't work. + tl.static_assert(BLOCK_M2 % BLOCK_N2 == 0) + step_n = BLOCK_N2 + + q_blk = tl.program_id(0) # Q-tile index + off_hz = tl.program_id(2) # fused (batch, head) + b = off_hz // H + h = off_hz % H + q_tiles = N_CTX // BLOCK_M2 + meta_base = ((b * H + h) * q_tiles + q_blk) + + kv_blocks = tl.load(q2k_num + meta_base) # int32 + kv_ptr = q2k_index + meta_base * max_kv_blks # ptr to list + block_size = tl.load(variable_block_sizes + q_blk) + + + for blk_idx in range(kv_blocks*2): + block_sparse_offset = (tl.load(kv_ptr + blk_idx//2).to(tl.int32)*2 + blk_idx%2) *step_n * stride_tok + kT = tl.load(kT_ptrs + block_sparse_offset) + vT = tl.load(vT_ptrs + block_sparse_offset) + qk = tl.dot(q, kT) + p = tl.math.exp2(qk - m) + mask = tl.arange(0, BLOCK_N2) < block_size.to(tl.int32) + p = tl.where(mask[None, :], p , 0.0) + # Compute dP and dS. + dp = tl.dot(do, vT).to(tl.float32) + ds = p * (dp - Di[:, None]) + ds = ds.to(tl.bfloat16) + # Compute dQ. + # NOTE: We need to de-scale dq in the end, because kT was pre-scaled. + dq += tl.dot(ds, tl.trans(kT)) + # Increment pointers. + return dq + + + +@triton.jit +def _attn_bwd(Q, K, V, sm_scale, # + DO, # + DQ, DK, DV, # + M, D, + q2k_index, q2k_num, max_kv_blks, + k2q_index, k2q_num, max_q_blks, + variable_block_sizes, + # shared by Q/K/V/DO. + stride_z, stride_h, stride_tok, stride_d, # + H, N_CTX, # + BLOCK_M1: tl.constexpr, # + BLOCK_N1: tl.constexpr, # + BLOCK_M2: tl.constexpr, # + BLOCK_N2: tl.constexpr, # + HEAD_DIM: tl.constexpr): + LN2 = 0.6931471824645996 # = ln(2) + + bhid = tl.program_id(2) + off_chz = (bhid * N_CTX).to(tl.int64) + adj = (stride_h * (bhid % H) + stride_z * (bhid // H)).to(tl.int64) + pid = tl.program_id(0) + + # offset pointers for batch/head + Q += adj + K += adj + V += adj + DO += adj + DQ += adj + DK += adj + DV += adj + M += off_chz + D += off_chz + + # load scales + offs_k = tl.arange(0, HEAD_DIM) + + start_n = pid * BLOCK_N1 + start_m = 0 + + offs_n = start_n + tl.arange(0, BLOCK_N1) + + dv = tl.zeros([BLOCK_N1, HEAD_DIM], dtype=tl.float32) + dk = tl.zeros([BLOCK_N1, HEAD_DIM], dtype=tl.float32) + + # load K and V: they stay in SRAM throughout the inner loop. + k = tl.load(K + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d) + v = tl.load(V + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d) + + + num_steps = N_CTX // BLOCK_M1 + + dk, dv = _attn_bwd_dkdv( # + dk, dv, # + Q, k, v, sm_scale, # + DO, # + M, D, # + k2q_index, k2q_num, max_q_blks, + variable_block_sizes, + stride_tok, stride_d, # + H, N_CTX, # + BLOCK_M1, BLOCK_N1, HEAD_DIM, # + start_n, start_m, num_steps # + ) + + dv_ptrs = DV + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d + tl.store(dv_ptrs, dv) + + # Write back dK. + dk *= sm_scale + dk_ptrs = DK + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d + tl.store(dk_ptrs, dk) + + # THIS BLOCK DOES DQ: + start_m = pid * BLOCK_M2 + end_n = 0 + + offs_m = start_m + tl.arange(0, BLOCK_M2) + + q = tl.load(Q + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d) + dq = tl.zeros([BLOCK_M2, HEAD_DIM], dtype=tl.float32) + do = tl.load(DO + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d) + + m = tl.load(M + offs_m) + m = m[:, None] + + num_steps = N_CTX // BLOCK_N2 + dq = _attn_bwd_dq(dq, q, K, V, # + do, m, D, # + q2k_index, q2k_num, max_kv_blks, + variable_block_sizes, + stride_tok, stride_d, # + H, N_CTX, # + BLOCK_M2, BLOCK_N2, HEAD_DIM, # + start_m, end_n, num_steps # + ) + # Write back dQ. + dq_ptrs = DQ + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d + dq *= LN2 + tl.store(dq_ptrs, dq) + + + + + +# ──────────────────────────── SPARSE ADDITION BEGIN ─────────────────────────── +def triton_block_sparse_attn_forward(q, k, v, q2k_index, q2k_num, variable_block_sizes): + B, H, T, D = q.shape + sm_scale = 1.0 / math.sqrt(D) + max_kv_blks = q2k_index.shape[-1] + assert T % 64 == 0, f"T must be a multiple of 64, but got {T}" + assert T // 64 == q2k_num.shape[-1], f"shape mismatch, T // 64 = {T // 64}, q2k_num.shape[-2] = {q2k_num.shape[-2]}" + o = torch.empty_like(q) + M = torch.empty((B, H, T), dtype=torch.float32, device=q.device) + + grid = lambda _: (triton.cdiv(T, 64), B * H, 1) + _attn_fwd_sparse[grid]( + q, k, v, sm_scale, + q2k_index, q2k_num, max_kv_blks, + variable_block_sizes, + M, o, + q.stride(0), q.stride(1), q.stride(2), q.stride(3), + k.stride(0), k.stride(1), k.stride(2), k.stride(3), + v.stride(0), v.stride(1), v.stride(2), v.stride(3), + o.stride(0), o.stride(1), o.stride(2), o.stride(3), + B, H, T, + HEAD_DIM=D, STAGE=3 + ) + + return o, M + +def triton_block_sparse_attn_backward(do, q, k, v, o, M, q2k_index, q2k_num, k2q_index, k2q_num, variable_block_sizes): + assert do.is_contiguous() + assert q.stride() == k.stride() == v.stride() == o.stride() == do.stride() + + B, H, T, D = q.shape + sm_scale = 1.0 / math.sqrt(D) + dq = torch.empty_like(q) + dk = torch.empty_like(k) + dv = torch.empty_like(v) + BATCH, N_HEAD, N_CTX = q.shape[:3] + BLOCK_M1, BLOCK_N1, BLOCK_M2, BLOCK_N2 = 32, 64, 64, 32 + RCP_LN2 = 1.4426950408889634 # = 1.0 / ln(2) + arg_k = k + arg_k = arg_k * (sm_scale * RCP_LN2) + PRE_BLOCK = 64 + assert N_CTX % PRE_BLOCK == 0 + pre_grid = (N_CTX // PRE_BLOCK, BATCH * N_HEAD) + delta = torch.empty_like(M) + _attn_bwd_preprocess[pre_grid]( + o, do, # + delta, # + BATCH, N_HEAD, N_CTX, # + BLOCK_M=PRE_BLOCK, HEAD_DIM=D # + ) + + + max_q_blks = k2q_index.shape[-1] + max_kv_blks = q2k_index.shape[-1] + + grid = (N_CTX // BLOCK_N1, 1, BATCH * N_HEAD) + _attn_bwd[grid]( + q, arg_k, v, sm_scale, do, dq, dk, dv, # + M, delta, # + q2k_index, q2k_num, max_kv_blks, + k2q_index, k2q_num, max_q_blks, + variable_block_sizes, + q.stride(0), q.stride(1), q.stride(2), q.stride(3), # + N_HEAD, N_CTX, # + BLOCK_M1=BLOCK_M1, BLOCK_N1=BLOCK_N1, # + BLOCK_M2=BLOCK_M2, BLOCK_N2=BLOCK_N2, # + HEAD_DIM=D # + ) + + return dq, dk, dv + + diff --git a/csrc/fastvideo_kernel/src/fastvideo_kernel/triton_kernels/index.py b/csrc/fastvideo_kernel/src/fastvideo_kernel/triton_kernels/index.py new file mode 100644 index 000000000..6225f15b4 --- /dev/null +++ b/csrc/fastvideo_kernel/src/fastvideo_kernel/triton_kernels/index.py @@ -0,0 +1,152 @@ + +## pytorch sdpa version of block sparse ## +import triton +import triton.language as tl +import torch + +@triton.jit +def topk_index_to_map_kernel( + map_ptr, + index_ptr, + map_bs_stride, + map_h_stride, + map_q_stride, + map_kv_stride, + index_bs_stride, + index_h_stride, + index_q_stride, + index_kv_stride, + topk, +): + b, h, q = tl.program_id(0), tl.program_id(1), tl.program_id(2) + index_ptr_base = index_ptr + b * index_bs_stride + h * index_h_stride + q * index_q_stride + map_ptr_base = map_ptr + b * map_bs_stride + h * map_h_stride + q * map_q_stride + + for i in tl.static_range(topk): + index = tl.load(index_ptr_base + i * index_kv_stride) + tl.store(map_ptr_base + index * map_kv_stride, 1.0) + +@triton.jit +def map_to_index_kernel( + map_ptr, + index_ptr, + index_num_ptr, + map_bs_stride, + map_h_stride, + map_q_stride, + map_kv_stride, + index_bs_stride, + index_h_stride, + index_q_stride, + index_kv_stride, + index_num_bs_stride, + index_num_h_stride, + index_num_q_stride, + num_kv_blocks, +): + b, h, q = tl.program_id(0), tl.program_id(1), tl.program_id(2) + index_ptr_base = index_ptr + b * index_bs_stride + h * index_h_stride + q * index_q_stride + map_ptr_base = map_ptr + b * map_bs_stride + h * map_h_stride + q * map_q_stride + + num = 0 + for i in tl.range(num_kv_blocks): + map_entry = tl.load(map_ptr_base + i * map_kv_stride) + if map_entry: + tl.store(index_ptr_base + num * index_kv_stride, i) + num += 1 + + tl.store( + index_num_ptr + b * index_num_bs_stride + h * index_num_h_stride + + q * index_num_q_stride, num) + +def topk_index_to_map(index: torch.Tensor, + num_kv_blocks: int, + transpose_map: bool = False): + """ + Convert topk indices to a map. + + Args: + index: [bs, h, num_q_blocks, topk] + The topk indices tensor. + num_kv_blocks: int + The number of key-value blocks in the block_map returned + transpose_map: bool + If True, the block_map will be transposed on the final two dimensions. + + Returns: + block_map: [bs, h, num_q_blocks, num_kv_blocks] + A binary map where 1 indicates that the q block attends to the kv block. + """ + bs, h, num_q_blocks, topk = index.shape + + if transpose_map is False: + block_map = torch.zeros((bs, h, num_q_blocks, num_kv_blocks), + dtype=torch.bool, + device=index.device) + else: + block_map = torch.zeros((bs, h, num_kv_blocks, num_q_blocks), + dtype=torch.bool, + device=index.device) + block_map = block_map.transpose(2, 3) + + grid = (bs, h, num_q_blocks) + topk_index_to_map_kernel[grid]( + block_map, + index, + block_map.stride(0), + block_map.stride(1), + block_map.stride(2), + block_map.stride(3), + index.stride(0), + index.stride(1), + index.stride(2), + index.stride(3), + topk=topk, + ) + + return block_map + +def map_to_index(block_map: torch.Tensor): + """ + Convert a block map to indices and counts. + + Args: + block_map: [bs, h, num_q_blocks, num_kv_blocks] + The block map tensor. + + Returns: + index: [bs, h, num_q_blocks, num_kv_blocks] + The indices of the blocks. + index_num: [bs, h, num_q_blocks] + The number of blocks for each q block. + """ + bs, h, num_q_blocks, num_kv_blocks = block_map.shape + + index = torch.full((block_map.shape), + -1, + dtype=torch.int32, + device=block_map.device) + index_num = torch.empty((bs, h, num_q_blocks), + dtype=torch.int32, + device=block_map.device) + + grid = (bs, h, num_q_blocks) + map_to_index_kernel[grid]( + block_map, + index, + index_num, + block_map.stride(0), + block_map.stride(1), + block_map.stride(2), + block_map.stride(3), + index.stride(0), + index.stride(1), + index.stride(2), + index.stride(3), + index_num.stride(0), + index_num.stride(1), + index_num.stride(2), + num_kv_blocks=num_kv_blocks, + ) + + return index, index_num diff --git a/csrc/fastvideo_kernel/src/fastvideo_kernel/vmoba.py b/csrc/fastvideo_kernel/src/fastvideo_kernel/vmoba.py new file mode 100644 index 000000000..6bac66ac6 --- /dev/null +++ b/csrc/fastvideo_kernel/src/fastvideo_kernel/vmoba.py @@ -0,0 +1,868 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapt from https://github.com/KwaiVGI/VMoBA/blob/main/src/vmoba.py + +import random +import time +import os +import torch +from typing import Tuple +try: + from flash_attn import flash_attn_varlen_func # Use the new flash attention function + from flash_attn.flash_attn_interface import _flash_attn_varlen_forward, _flash_attn_varlen_backward +except ImportError: + def _unsupported(*args, **kwargs): + raise ImportError("flash-attn is not installed. Please install it, e.g., `pip install flash-attn`.") + _flash_attn_varlen_forward = _unsupported + _flash_attn_varlen_backward = _unsupported + flash_attn_varlen_func = _unsupported + +from functools import lru_cache +from einops import rearrange + + +@lru_cache(maxsize=16) +def calc_chunks(cu_seqlen, moba_chunk_size): + """ + Calculate chunk boundaries. + + For vision tasks we include all chunks (even the last one which might be shorter) + so that every chunk can be selected. + """ + batch_sizes = cu_seqlen[1:] - cu_seqlen[:-1] + batch_num_chunk = (batch_sizes + (moba_chunk_size - 1)) // moba_chunk_size + cu_num_chunk = torch.ones( + batch_num_chunk.numel() + 1, + device=cu_seqlen.device, + dtype=batch_num_chunk.dtype, + ) + cu_num_chunk[1:] = batch_num_chunk.cumsum(dim=0) + num_chunk = cu_num_chunk[-1] + chunk_sizes = torch.full( + (num_chunk + 1,), moba_chunk_size, dtype=torch.int32, device=cu_seqlen.device + ) + chunk_sizes[0] = 0 + batch_last_chunk_size = batch_sizes - (batch_num_chunk - 1) * moba_chunk_size + chunk_sizes[cu_num_chunk[1:]] = batch_last_chunk_size + cu_chunk = chunk_sizes.cumsum(dim=-1, dtype=torch.int32) + chunk_to_batch = torch.zeros( + (num_chunk,), dtype=torch.int32, device=cu_seqlen.device + ) + chunk_to_batch[cu_num_chunk[1:-1]] = 1 + chunk_to_batch = chunk_to_batch.cumsum(dim=0, dtype=torch.int32) + + # Do not filter out any chunk + filtered_chunk_indices = torch.arange( + num_chunk, device=cu_seqlen.device, dtype=torch.int32 + ) + num_filtered_chunk = num_chunk + + return cu_chunk, filtered_chunk_indices, num_filtered_chunk, chunk_to_batch + + +# --- Threshold Selection Helper Functions --- + +def _select_threshold_query_head( + gate: torch.Tensor, + valid_gate_mask: torch.Tensor, + gate_self_chunk_mask: torch.Tensor, + simsum_threshold: float +) -> torch.Tensor: + """ + Selects chunks for each pair based on threshold. + Normalization and sorting happen along the chunk dimension (dim=0). + """ + C, H, S = gate.shape + eps = 1e-6 + + # LSE‐style normalization per (across chunks) + gate_masked = torch.where(valid_gate_mask, gate, -torch.inf) # Use -inf for max + gate_min_val = torch.where(valid_gate_mask, gate, torch.inf) # Use +inf for min + + row_min = gate_min_val.amin(dim=0) # (H, S) + row_max = gate_masked.amax(dim=0) # (H, S) + denom = row_max - row_min + denom = torch.where(denom <= eps, torch.ones_like(denom), denom) # avoid divide‑by‑zero + + gate_norm = (gate - row_min.unsqueeze(0)) / denom.unsqueeze(0) + gate_norm = torch.where(valid_gate_mask, gate_norm, 0.0) # (C, H, S) + + # 1) pull out the self‐chunk’s normalized weight for each + self_norm = (gate_norm * gate_self_chunk_mask).sum(dim=0) # (H, S) + + # 2) compute how much more normalized weight we need beyond self + total_norm_sum = gate_norm.sum(dim=0) # (H, S) + remain_ratio = simsum_threshold - self_norm / (total_norm_sum + eps) # (H, S) + remain_ratio = torch.clamp(remain_ratio, min=0.0) # if already ≥ thresh, no extra needed + + # 3) zero out the self‐chunk in a copy, so we only sort “others” + others_norm = gate_norm.clone() + others_norm[gate_self_chunk_mask] = 0.0 + + # 4) sort the other chunks by descending norm, per + sorted_norm, sorted_idx = torch.sort(others_norm, descending=True, dim=0) # (C, H, S) + + # 5) cumulative‑sum the sorted norms per + cumsum_others = sorted_norm.cumsum(dim=0) # (C, H, S) + + # 6) for each , find the smallest k where cumsum_ratio ≥ remain_ratio + ratio = cumsum_others / (total_norm_sum.unsqueeze(0) + eps) # (C, H, S) + cond = ratio >= remain_ratio.unsqueeze(0) # (C, H, S) boolean mask + any_cond = cond.any(dim=0) # (H, S) + # Find the index of the first True value along dim 0. If none, use C-1. + cutoff = torch.where(any_cond, cond.float().argmax(dim=0), torch.full_like(any_cond, fill_value=C - 1)) # (H, S) + + # 7) build a mask in sorted order up to that cutoff + idx_range = torch.arange(C, device=gate.device).view(-1, 1, 1) # (C, 1, 1) + sorted_mask = idx_range <= cutoff.unsqueeze(0) # (C, H, S) + + # 8) scatter it back to original chunk order + others_mask = torch.zeros_like(gate, dtype=torch.bool) + others_mask.scatter_(0, sorted_idx, sorted_mask) + + # 9) finally, include every self‐chunk plus all selected others + final_gate_mask = valid_gate_mask & (others_mask | gate_self_chunk_mask) + + return final_gate_mask + + +def _select_threshold_block( + gate: torch.Tensor, + valid_gate_mask: torch.Tensor, + gate_self_chunk_mask: torch.Tensor, + simsum_threshold: float +) -> torch.Tensor: + """ + Selects pairs for each block based on threshold. + Normalization and sorting happen across the head and sequence dimensions (dim=1, 2). + """ + C, H, S = gate.shape + HS = H * S + eps = 1e-6 + + # LSE‐style normalization per block (across heads and queries) + gate_masked = torch.where(valid_gate_mask, gate, -torch.inf) # Use -inf for max + gate_min_val = torch.where(valid_gate_mask, gate, torch.inf) # Use +inf for min + + block_max = gate_masked.amax(dim=(1, 2), keepdim=True) # (C, 1, 1) + block_min = gate_min_val.amin(dim=(1, 2), keepdim=True) # (C, 1, 1) + block_denom = block_max - block_min + block_denom = torch.where(block_denom <= eps, torch.ones_like(block_denom), block_denom) # (C, 1, 1) + + gate_norm = (gate - block_min) / block_denom # (C, H, S) + gate_norm = torch.where(valid_gate_mask, gate_norm, 0.0) # (C, H, S) + + # 1) identify normalized weights of entries that *are* self-chunks (from query perspective) + self_norm_entries = gate_norm * gate_self_chunk_mask # (C, H, S) + # Sum these weights *per block* + self_norm_sum_per_block = self_norm_entries.sum(dim=(1, 2)) # (C,) + + # 2) compute how much more normalized weight each block needs beyond its self-chunk contributions + total_norm_sum_per_block = gate_norm.sum(dim=(1, 2)) # (C,) + remain_ratio = simsum_threshold - self_norm_sum_per_block / (total_norm_sum_per_block + eps) # (C,) + remain_ratio = torch.clamp(remain_ratio, min=0.0) # (C,) + + # 3) zero out the self‐chunk entries in a copy, so we only sort “others” + others_norm = gate_norm.clone() + others_norm[gate_self_chunk_mask] = 0.0 # Zero out self entries + + # 4) sort the other pairs by descending norm, per block + others_flat = others_norm.contiguous().view(C, HS) # (C, H*S) + sorted_others_flat, sorted_indices_flat = torch.sort(others_flat, dim=1, descending=True) # (C, H*S) + + # 5) cumulative‑sum the sorted norms per block + cumsum_others_flat = sorted_others_flat.cumsum(dim=1) # (C, H*S) + + # 6) for each block, find the smallest k where cumsum_ratio ≥ remain_ratio + ratio_flat = cumsum_others_flat / (total_norm_sum_per_block.unsqueeze(1) + eps) # (C, H*S) + cond_flat = ratio_flat >= remain_ratio.unsqueeze(1) # (C, H*S) boolean mask + any_cond = cond_flat.any(dim=1) # (C,) + # Find the index of the first True value along dim 1. If none, use HS-1. + cutoff_flat = torch.where(any_cond, cond_flat.float().argmax(dim=1), torch.full_like(any_cond, fill_value=HS - 1)) # (C,) + + # 7) build a mask in sorted order up to that cutoff per block + idx_range_flat = torch.arange(HS, device=gate.device).unsqueeze(0) # (1, H*S) + sorted_mask_flat = idx_range_flat <= cutoff_flat.unsqueeze(1) # (C, H*S) + + # 8) scatter it back to original order per block + others_mask_flat = torch.zeros_like(others_flat, dtype=torch.bool) # (C, H*S) + others_mask_flat.scatter_(1, sorted_indices_flat, sorted_mask_flat) + others_mask = others_mask_flat.view(C, H, S) # (C, H, S) + + # 9) finally, include every self‐chunk entry plus all selected others + final_gate_mask = valid_gate_mask & (others_mask | gate_self_chunk_mask) + + return final_gate_mask + + +def _select_threshold_overall( + gate: torch.Tensor, + valid_gate_mask: torch.Tensor, + gate_self_chunk_mask: torch.Tensor, + simsum_threshold: float +) -> torch.Tensor: + """ + Selects triplets globally based on threshold. + Normalization and sorting happen across all valid entries. + """ + C, H, S = gate.shape + CHS = C * H * S + eps = 1e-6 + + # LSE‐style normalization globally across all valid entries + gate_masked = torch.where(valid_gate_mask, gate, -torch.inf) # Use -inf for max + gate_min_val = torch.where(valid_gate_mask, gate, torch.inf) # Use +inf for min + + overall_max = gate_masked.max() # scalar + overall_min = gate_min_val.min() # scalar + overall_denom = overall_max - overall_min + overall_denom = torch.where(overall_denom <= eps, torch.tensor(1.0, device=gate.device, dtype=gate.dtype), overall_denom) + + gate_norm = (gate - overall_min) / overall_denom # (C, H, S) + gate_norm = torch.where(valid_gate_mask, gate_norm, 0.0) # (C, H, S) + + # 1) identify normalized weights of entries that *are* self-chunks + self_norm_entries = gate_norm * gate_self_chunk_mask # (C, H, S) + # Sum these weights globally + self_norm_sum_overall = self_norm_entries.sum() # scalar + + # 2) compute how much more normalized weight is needed globally beyond self-chunk contributions + total_norm_sum_overall = gate_norm.sum() # scalar + remain_ratio = simsum_threshold - self_norm_sum_overall / (total_norm_sum_overall + eps) # scalar + remain_ratio = torch.clamp(remain_ratio, min=0.0) # scalar + + # 3) zero out the self‐chunk entries in a copy, so we only sort “others” + others_norm = gate_norm.clone() + others_norm[gate_self_chunk_mask] = 0.0 # Zero out self entries + + # 4) sort all other entries by descending norm, globally + others_flat = others_norm.flatten() # (C*H*S,) + valid_others_mask_flat = valid_gate_mask.flatten() & ~gate_self_chunk_mask.flatten() # Mask for valid, non-self entries + + # Only sort the valid 'other' entries + valid_others_indices = torch.where(valid_others_mask_flat)[0] + valid_others_values = others_flat[valid_others_indices] + + sorted_others_values, sort_perm = torch.sort(valid_others_values, descending=True) # (N_valid_others,) + sorted_original_indices = valid_others_indices[sort_perm] # Original indices in C*H*S space, sorted by value + + # 5) cumulative‑sum the sorted valid 'other' norms globally + cumsum_others_values = sorted_others_values.cumsum(dim=0) # (N_valid_others,) + + # 6) find the smallest k where cumsum_ratio ≥ remain_ratio globally + ratio_values = cumsum_others_values / (total_norm_sum_overall + eps) # (N_valid_others,) + cond_values = ratio_values >= remain_ratio # (N_valid_others,) boolean mask + any_cond = cond_values.any() # scalar + + # Find the index of the first True value in the *sorted* list. If none, use all valid others. + cutoff_idx_in_sorted = torch.where( + any_cond, + cond_values.float().argmax(dim=0), + torch.tensor(len(sorted_others_values) - 1, device=gate.device, dtype=torch.long) + ) + + # 7) build a mask selecting the top-k others based on the cutoff + # Select the original indices corresponding to the top entries in the sorted list + selected_other_indices = sorted_original_indices[:cutoff_idx_in_sorted + 1] + + # 8) create the mask in the original flat shape + others_mask_flat = torch.zeros_like(others_flat, dtype=torch.bool) # (C*H*S,) + if selected_other_indices.numel() > 0: # Check if any 'other' indices were selected + others_mask_flat[selected_other_indices] = True + others_mask = others_mask_flat.view(C, H, S) # (C, H, S) + + # 9) finally, include every self‐chunk entry plus all selected others + final_gate_mask = valid_gate_mask & (others_mask | gate_self_chunk_mask) + + return final_gate_mask + + +def _select_threshold_head_global( + gate: torch.Tensor, + valid_gate_mask: torch.Tensor, + gate_self_chunk_mask: torch.Tensor, + simsum_threshold: float +) -> torch.Tensor: + """ + Selects globally for each head based on threshold. + """ + C, H, S = gate.shape + eps = 1e-6 + + # 1) LSE‐style normalization per head (across chunks and sequence dims) + gate_masked = torch.where(valid_gate_mask, gate, -torch.inf) + gate_min_val = torch.where(valid_gate_mask, gate, torch.inf) + + max_per_head = gate_masked.amax(dim=(0, 2), keepdim=True) # (1, H, 1) + min_per_head = gate_min_val.amin(dim=(0, 2), keepdim=True) # (1, H, 1) + denom = max_per_head - min_per_head + denom = torch.where(denom <= eps, torch.ones_like(denom), denom) + + gate_norm = (gate - min_per_head) / denom + gate_norm = torch.where(valid_gate_mask, gate_norm, 0.0) # (C, H, S) + + # 2) sum normalized self‐chunk contributions per head + self_norm_sum = (gate_norm * gate_self_chunk_mask).sum(dim=(0, 2)) # (H,) + + # 3) total normalized sum per head + total_norm_sum = gate_norm.sum(dim=(0, 2)) # (H,) + + # 4) how much more normalized weight needed per head + remain_ratio = simsum_threshold - self_norm_sum / (total_norm_sum + eps) # (H,) + remain_ratio = torch.clamp(remain_ratio, min=0.0) + + # 5) zero out self‐chunk entries to focus on "others" + others_norm = gate_norm.clone() + others_norm[gate_self_chunk_mask] = 0.0 # (C, H, S) + + # 6) flatten chunk and sequence dims, per head + CS = C * S + others_flat = others_norm.permute(1, 0, 2).reshape(H, CS) # (H, C*S) + valid_flat = (valid_gate_mask & ~gate_self_chunk_mask) \ + .permute(1, 0, 2).reshape(H, CS) # (H, C*S) + + # 7) vectorized selection of “others” per head + masked_flat = torch.where(valid_flat, others_flat, torch.zeros_like(others_flat)) + sorted_vals, sorted_idx = torch.sort(masked_flat, dim=1, descending=True) # (H, C*S) + + cumsum_vals = sorted_vals.cumsum(dim=1) # (H, C*S) + ratio_vals = cumsum_vals / (total_norm_sum.unsqueeze(1) + eps) # (H, C*S) + cond = ratio_vals >= remain_ratio.unsqueeze(1) # (H, C*S) + + has_cutoff = cond.any(dim=1) # (H,) + default = torch.full((H,), CS - 1, device=gate.device, dtype=torch.long) + cutoff = torch.where(has_cutoff, cond.float().argmax(dim=1), default) # (H,) + + idx_range = torch.arange(CS, device=gate.device).unsqueeze(0) # (1, C*S) + sorted_mask = idx_range <= cutoff.unsqueeze(1) # (H, C*S) + + selected_flat = torch.zeros_like(valid_flat) # (H, C*S) + selected_flat.scatter_(1, sorted_idx, sorted_mask) # (H, C*S) + + # 8) reshape selection mask back to (C, H, S) + others_mask = selected_flat.reshape(H, C, S).permute(1, 0, 2) # (C, H, S) + + # 9) include self‐chunks plus selected others, and obey valid mask + final_gate_mask = valid_gate_mask & (gate_self_chunk_mask | others_mask) + + return final_gate_mask + + +class MixedAttention(torch.autograd.Function): + @staticmethod + def forward( + ctx, + q, + k, + v, + self_attn_cu_seqlen, + moba_q, + moba_kv, + moba_cu_seqlen_q, + moba_cu_seqlen_kv, + max_seqlen, + moba_chunk_size, + moba_q_sh_indices, + ): + ctx.max_seqlen = max_seqlen + ctx.moba_chunk_size = moba_chunk_size + ctx.softmax_scale = softmax_scale = q.shape[-1] ** (-0.5) + + # Non-causal self-attention branch + # return out, softmax_lse, S_dmask, rng_state + self_attn_out_sh, self_attn_lse_hs, _, _ = _flash_attn_varlen_forward( + q=q, + k=k, + v=v, + cu_seqlens_q=self_attn_cu_seqlen, + cu_seqlens_k=self_attn_cu_seqlen, + max_seqlen_q=max_seqlen, + max_seqlen_k=max_seqlen, + softmax_scale=softmax_scale, + causal=False, + dropout_p=0.0, + ) + # MOBA attention branch (non-causal) + moba_attn_out, moba_attn_lse_hs, _, _ = _flash_attn_varlen_forward( + q=moba_q, + k=moba_kv[:, 0], + v=moba_kv[:, 1], + cu_seqlens_q=moba_cu_seqlen_q, + cu_seqlens_k=moba_cu_seqlen_kv, + max_seqlen_q=max_seqlen, + max_seqlen_k=moba_chunk_size, + softmax_scale=softmax_scale, + causal=False, + dropout_p=0.0, + ) + + self_attn_lse_sh = self_attn_lse_hs.t().contiguous() + moba_attn_lse = moba_attn_lse_hs.t().contiguous() + + output = torch.zeros((q.shape[0], q.shape[1], q.shape[2]), device=q.device, dtype=torch.float32) + output_2d = output.view(-1, q.shape[2]) + + max_lse_1d = self_attn_lse_sh.view(-1) + max_lse_1d = max_lse_1d.index_reduce( + 0, moba_q_sh_indices, moba_attn_lse.view(-1), "amax" + ) + self_attn_lse_sh = self_attn_lse_sh - max_lse_1d.view_as(self_attn_lse_sh) + moba_attn_lse = ( + moba_attn_lse.view(-1) + .sub(max_lse_1d.index_select(0, moba_q_sh_indices)) + .reshape_as(moba_attn_lse) + ) + + mixed_attn_se_sh = self_attn_lse_sh.exp() + moba_attn_se = moba_attn_lse.exp() + + mixed_attn_se_sh.view(-1).index_add_( + 0, moba_q_sh_indices, moba_attn_se.view(-1) + ) + mixed_attn_lse_sh = mixed_attn_se_sh.log() + + # Combine self-attention output + factor = (self_attn_lse_sh - mixed_attn_lse_sh).exp() # [S, H] + self_attn_out_sh = self_attn_out_sh * factor.unsqueeze(-1) + output_2d += self_attn_out_sh.reshape_as(output_2d) + + # Combine MOBA attention output + mixed_attn_lse = ( + mixed_attn_lse_sh.view(-1) + .index_select(0, moba_q_sh_indices) + .view_as(moba_attn_lse) + ) + factor = (moba_attn_lse - mixed_attn_lse).exp() # [S, H] + moba_attn_out = moba_attn_out * factor.unsqueeze(-1) + raw_attn_out = moba_attn_out.view(-1, moba_attn_out.shape[-1]) + output_2d.index_add_(0, moba_q_sh_indices, raw_attn_out) + output = output.to(q.dtype) + mixed_attn_lse_sh = mixed_attn_lse_sh + max_lse_1d.view_as(mixed_attn_se_sh) + ctx.save_for_backward( + output, + mixed_attn_lse_sh, + q, + k, + v, + self_attn_cu_seqlen, + moba_q, + moba_kv, + moba_cu_seqlen_q, + moba_cu_seqlen_kv, + moba_q_sh_indices, + ) + + return output + + @staticmethod + def backward(ctx, d_output): + + max_seqlen = ctx.max_seqlen + moba_chunk_size = ctx.moba_chunk_size + softmax_scale = ctx.softmax_scale + + ( + output, + mixed_attn_vlse_sh, + q, + k, + v, + self_attn_cu_seqlen, + moba_q, + moba_kv, + moba_cu_seqlen_q, + moba_cu_seqlen_kv, + moba_q_sh_indices, + ) = ctx.saved_tensors + + d_output = d_output.contiguous() + + dq = torch.empty_like(q) + dk = torch.empty_like(k) + dv = torch.empty_like(v) + _ = _flash_attn_varlen_backward( + dout=d_output, + q=q, + k=k, + v=v, + out=output, + softmax_lse=mixed_attn_vlse_sh.t().contiguous(), + dq=dq, + dk=dk, + dv=dv, + cu_seqlens_q=self_attn_cu_seqlen, + cu_seqlens_k=self_attn_cu_seqlen, + max_seqlen_q=max_seqlen, + max_seqlen_k=max_seqlen, + softmax_scale=softmax_scale, + causal=False, + dropout_p=0.0, + softcap=0.0, + alibi_slopes=None, + deterministic=True, + window_size_left=-1, + window_size_right=-1 + ) + + headdim = q.shape[-1] + d_moba_output = ( + d_output.view(-1, headdim).index_select(0, moba_q_sh_indices).unsqueeze(1) + ) + moba_output = ( + output.view(-1, headdim).index_select(0, moba_q_sh_indices).unsqueeze(1) + ) + + mixed_attn_vlse = ( + mixed_attn_vlse_sh.view(-1).index_select(0, moba_q_sh_indices).view(1, -1) + ) + + dmq = torch.empty_like(moba_q) + dmkv = torch.empty_like(moba_kv) + _ = _flash_attn_varlen_backward( + dout=d_moba_output, + q=moba_q, + k=moba_kv[:, 0], + v=moba_kv[:, 1], + out=moba_output, + softmax_lse=mixed_attn_vlse, + dq=dmq, + dk=dmkv[:,0], + dv=dmkv[:,1], + cu_seqlens_q=moba_cu_seqlen_q, + cu_seqlens_k=moba_cu_seqlen_kv, + max_seqlen_q=max_seqlen, + max_seqlen_k=moba_chunk_size, + softmax_scale=softmax_scale, + causal=False, + dropout_p=0.0, + softcap=0.0, + alibi_slopes=None, + deterministic=True, + window_size_left=-1, + window_size_right=-1 + ) + + return dq, dk, dv, None, dmq, dmkv, None, None, None, None, None + + +def moba_attn_varlen( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + cu_seqlens: torch.Tensor, + max_seqlen: int, + moba_chunk_size: int, + moba_topk: int, + select_mode: str = 'threshold', # "topk" or "threshold" + simsum_threshold: float = 0.25, + threshold_type: str = 'query_head', +) -> torch.Tensor: + """ + Accelerated MOBA attention for vision tasks with proper LSE normalization. + + This version: + - Splits KV into chunks. + - For each query head, selects the top-k relevant KV chunks (including the self chunk) + by amplifying the diagonal (self-chunk) logits. + - Aggregates the attention outputs from the selected chunks using a log-sum-exp + reduction so that attending to each query over the selected chunks is equivalent + to the original algorithm. + """ + # Stack keys and values. + kv = torch.stack((k, v), dim=1) + seqlen, num_head, head_dim = q.shape + + # Compute chunk boundaries. + cu_chunk, filtered_chunk_indices, num_filtered_chunk, chunk_to_batch = calc_chunks( + cu_seqlens, moba_chunk_size + ) + + self_attn_cu_seqlen = cu_chunk + + # Update top-k selection to include the self chunk. + moba_topk = min(moba_topk, num_filtered_chunk) + + # --- Build filtered KV from chunks --- + chunk_starts = cu_chunk[filtered_chunk_indices] # [num_filtered_chunk] + chunk_ends = cu_chunk[filtered_chunk_indices + 1] # [num_filtered_chunk] + chunk_lengths = chunk_ends - chunk_starts # [num_filtered_chunk] + max_chunk_len = int(chunk_lengths.max().item()) + + range_tensor = torch.arange(max_chunk_len, device=kv.device, dtype=chunk_starts.dtype).unsqueeze(0) + indices = chunk_starts.unsqueeze(1) + range_tensor + indices = torch.clamp(indices, max=kv.shape[0] - 1) + valid_mask = range_tensor < chunk_lengths.unsqueeze(1) + gathered = kv[indices.view(-1)].view(num_filtered_chunk, max_chunk_len, *kv.shape[1:]) + gathered = gathered * valid_mask.unsqueeze(-1).unsqueeze(-1).unsqueeze(-1).type_as(gathered) + + # Compute key_gate_weight over valid tokens. + key_values = gathered[:, :, 0].float() # [num_filtered_chunk, max_chunk_len, num_head, head_dim] + valid_mask_exp = valid_mask.unsqueeze(-1).unsqueeze(-1) + key_sum = (key_values * valid_mask_exp).sum(dim=1) + divisor = valid_mask.sum(dim=1).unsqueeze(-1).unsqueeze(-1) + key_gate_weight = key_sum / divisor # [num_filtered_chunk, num_head, head_dim] + + # Compute gate logits between key_gate_weight and queries. + q_float = q.float() + # gate = torch.einsum("nhd,shd->nhs", key_gate_weight, q_float) # [num_filtered_chunk, num_head, seqlen] + gate = torch.bmm(key_gate_weight.permute(1, 0, 2), q_float.permute(1, 0, 2).transpose(1, 2)).permute(1, 0, 2) + + # Amplify the diagonal (self chunk) contributions. + gate_seq_idx = torch.arange(seqlen, device=q.device, dtype=torch.int32).unsqueeze(0).expand(num_filtered_chunk, seqlen) + chunk_start = cu_chunk[filtered_chunk_indices] # [num_filtered_chunk] + chunk_end = cu_chunk[filtered_chunk_indices + 1] # [num_filtered_chunk] + gate_self_chunk_mask = ((gate_seq_idx >= chunk_start.unsqueeze(1)) & + (gate_seq_idx < chunk_end.unsqueeze(1))).unsqueeze(1).expand(-1, num_head, -1) + amplification_factor = 1e9 # Example factor; adjust as needed. + origin_gate = gate.clone() + gate = gate.clone() + if select_mode == "topk": + gate[gate_self_chunk_mask] += amplification_factor + + # Exclude positions that are outside the valid batch boundaries. + batch_starts = cu_seqlens[chunk_to_batch[filtered_chunk_indices]] + batch_ends = cu_seqlens[chunk_to_batch[filtered_chunk_indices] + 1] + gate_batch_start_mask = gate_seq_idx < batch_starts.unsqueeze(1) + gate_batch_end_mask = gate_seq_idx >= batch_ends.unsqueeze(1) + gate_inf_mask = gate_batch_start_mask | gate_batch_end_mask + gate.masked_fill_(gate_inf_mask.unsqueeze(1), -float("inf")) + + if select_mode == 'topk': + # We amplify self‐chunk in gate already, so self entries will rank highest. + valid_gate_mask = gate != -float("inf") + if threshold_type == 'query_head': + # === per‐ top-k across chunks (original behavior) === + # gate: (C, H, S) + _, gate_topk_idx = torch.topk(gate, k=moba_topk, dim=0, largest=True, sorted=False) + gate_idx_mask = torch.zeros_like(gate, dtype=torch.bool) + gate_idx_mask.scatter_(0, gate_topk_idx, True) + gate_mask = valid_gate_mask & gate_idx_mask + elif threshold_type == 'overall': + # === global top-k across all (chunk, head, seq) entries === + C, H, S = gate.shape + flat_gate = gate.flatten() + flat_mask = valid_gate_mask.flatten() + flat_gate_masked = torch.where(flat_mask, flat_gate, -float("inf")) + # pick topk global entries + vals, idx = torch.topk(flat_gate_masked, k=moba_topk * H * S, largest=True, sorted=False) + others_mask_flat = torch.zeros_like(flat_mask, dtype=torch.bool) + others_mask_flat[idx] = True + gate_mask = (valid_gate_mask.flatten() & others_mask_flat).view(gate.shape) + elif threshold_type == 'head_global': + # per-head top-k across all chunks and sequence positions + C, H, S = gate.shape + CS = C * S + flat_gate = gate.permute(1, 0, 2).reshape(H, CS) + flat_valid = valid_gate_mask.permute(1, 0, 2).reshape(H, CS) + flat_gate_masked = torch.where(flat_valid, flat_gate, torch.full_like(flat_gate, -float('inf'))) + # pick top-k indices per head + _, topk_idx = torch.topk(flat_gate_masked, k=moba_topk * S, dim=1, largest=True, sorted=False) + gate_idx_flat = torch.zeros_like(flat_valid, dtype=torch.bool) + gate_idx_flat.scatter_(1, topk_idx, True) + gate_mask = gate_idx_flat.reshape(H, C, S).permute(1, 0, 2) + else: + raise ValueError( + f"Invalid threshold_type for topk: {threshold_type}. " + "Choose 'query_head', 'block', or 'overall'." + ) + elif select_mode == 'threshold': + # Delegate to the specific thresholding function + valid_gate_mask = gate != -float("inf") # (num_chunk, num_head, seqlen) + if threshold_type == 'query_head': + gate_mask = _select_threshold_query_head(gate, valid_gate_mask, gate_self_chunk_mask, simsum_threshold) + elif threshold_type == 'block': + gate_mask = _select_threshold_block(gate, valid_gate_mask, gate_self_chunk_mask, simsum_threshold) + elif threshold_type == 'overall': + gate_mask = _select_threshold_overall(gate, valid_gate_mask, gate_self_chunk_mask, simsum_threshold) + elif threshold_type == 'head_global': + gate_mask = _select_threshold_head_global(gate, valid_gate_mask, gate_self_chunk_mask, simsum_threshold) + else: + raise ValueError(f"Invalid threshold_type: {threshold_type}. Choose 'query_head', 'block', or 'overall'.") + else: + raise ValueError(f"Invalid select_mode: {select_mode}. Choose 'topk' or 'threshold'.") + + # eliminate self_chunk in MoBA branch + gate_mask = gate_mask & ~gate_self_chunk_mask + # if gate_mask is all false, perform flash_attn instead + if gate_mask.sum() == 0: + return flash_attn_varlen_func( + q, k, v, cu_seqlens, cu_seqlens, max_seqlen, max_seqlen, causal=False + ) + + # Determine which query positions are selected. + # nonzero_indices has shape [N, 3] where each row is [chunk_index, head_index, seq_index]. + moba_q_indices = gate_mask.reshape(gate_mask.shape[0], -1).nonzero(as_tuple=True)[-1] # [(h s k)] + moba_q_sh_indices = (moba_q_indices % seqlen) * num_head + (moba_q_indices // seqlen) + moba_q = rearrange(q, "s h d -> (h s) d").index_select(0, moba_q_indices).unsqueeze(1) + + # Build cumulative sequence lengths for the selected queries. + moba_seqlen_q = gate_mask.sum(dim=-1).flatten() + q_zero_mask = moba_seqlen_q == 0 + valid_expert_mask = ~q_zero_mask + if q_zero_mask.sum() > 0: + moba_seqlen_q = moba_seqlen_q[valid_expert_mask] + moba_cu_seqlen_q = torch.cat( + ( + torch.tensor([0], device=q.device, dtype=moba_seqlen_q.dtype), + moba_seqlen_q.cumsum(dim=0), + ), + dim=0, + ).to(torch.int32) + + # Rearrange gathered KV for the MOBA branch. + experts_tensor = rearrange(gathered, "nc cl two h d -> (nc h) cl two d") + valid_expert_lengths = chunk_lengths.unsqueeze(1).expand(num_filtered_chunk, num_head).reshape(-1).to(torch.int32) + if q_zero_mask.sum() > 0: + experts_tensor = experts_tensor[valid_expert_mask] + valid_expert_lengths = valid_expert_lengths[valid_expert_mask] + + seq_range = torch.arange(experts_tensor.shape[1], device=experts_tensor.device).unsqueeze(0) + mask = seq_range < valid_expert_lengths.unsqueeze(1) + moba_kv = experts_tensor[mask] # Shape: ((nc h cl_valid) two d) + moba_kv = moba_kv.unsqueeze(2) # Shape: ((nc h cl_valid) two 1 d) + + moba_cu_seqlen_kv = torch.cat( + [torch.zeros(1, device=experts_tensor.device, dtype=torch.int32), + valid_expert_lengths.cumsum(dim=0)], + dim=0, + ).to(torch.int32) + + assert ( + moba_cu_seqlen_kv.shape == moba_cu_seqlen_q.shape + ), f"Mismatch between moba_cu_seqlen_kv.shape and moba_cu_seqlen_q.shape: {moba_cu_seqlen_kv.shape} vs {moba_cu_seqlen_q.shape}" + + return MixedAttention.apply( + q, + k, + v, + self_attn_cu_seqlen, + moba_q, + moba_kv, + moba_cu_seqlen_q, + moba_cu_seqlen_kv, + max_seqlen, + moba_chunk_size, + moba_q_sh_indices, + ) + + +def process_moba_input( + x, + patch_resolution, + chunk_size, +): + """ + Process inputs for the attention function. + + Args: + x (torch.Tensor): Input tensor with shape [batch_size, num_patches, num_heads, head_dim]. + patch_resolution (tuple): Tuple containing the patch resolution (t, h, w). + chunk_size (int): Size of the chunk. (maybe tuple or int, according to chunk type) + + Returns: + torch.Tensor: Processed input tensor. + """ + if isinstance(chunk_size, float) or isinstance(chunk_size, int): + moba_chunk_size = int(chunk_size * patch_resolution[1] * patch_resolution[2]) + else: + assert isinstance(chunk_size, (Tuple, list)), f"chunk_size should be a tuple, list, or int, now it is: {type(chunk_size)}" + if len(chunk_size) == 2: + assert patch_resolution[1] % chunk_size[0] == 0 and patch_resolution[2] % chunk_size[1] == 0, f"spatial patch_resolution {patch_resolution[1:]} should be divisible by 2d chunk_size {chunk_size}" + nch, ncw = patch_resolution[1] // chunk_size[0], patch_resolution[2] // chunk_size[1] + x = rearrange(x, "b (t nch ch ncw cw) n d -> b (nch ncw t ch cw) n d", t=patch_resolution[0], nch=nch, ncw=ncw, ch=chunk_size[0], cw=chunk_size[1]) + moba_chunk_size = patch_resolution[0] * chunk_size[0] * chunk_size[1] + elif len(chunk_size) == 3: + assert patch_resolution[0] % chunk_size[0] == 0 and patch_resolution[1] % chunk_size[1] == 0 and patch_resolution[2] % chunk_size[2] == 0, f"patch_resolution {patch_resolution} should be divisible by 3d chunk_size {chunk_size}" + nct, nch, ncw = patch_resolution[0] // chunk_size[0], patch_resolution[1] // chunk_size[1], patch_resolution[2] // chunk_size[2] + x = rearrange(x, "b (nct ct nch ch ncw cw) n d -> b (nct nch ncw ct ch cw) n d", nct=nct, nch=nch, ncw=ncw, ct=chunk_size[0], ch=chunk_size[1], cw=chunk_size[2]) + moba_chunk_size = chunk_size[0] * chunk_size[1] * chunk_size[2] + else: + raise ValueError(f"chunk_size should be a int, or a tuple of length 2 or 3, now it is: {len(chunk_size)}") + + return x, moba_chunk_size + + +def process_moba_output( + x, + patch_resolution, + chunk_size, +): + if isinstance(chunk_size, float) or isinstance(chunk_size, int): + pass + elif len(chunk_size) == 2: + x = rearrange(x, "b (nch ncw t ch cw) n d -> b (t nch ch ncw cw) n d", nch=patch_resolution[1] // chunk_size[0], ncw=patch_resolution[2] // chunk_size[1], t=patch_resolution[0], ch=chunk_size[0], cw=chunk_size[1]) + elif len(chunk_size) == 3: + x = rearrange(x, "b (nct nch ncw ct ch cw) n d -> b (nct ct nch ch ncw cw) n d", nct=patch_resolution[0] // chunk_size[0], nch=patch_resolution[1] // chunk_size[1], ncw=patch_resolution[2] // chunk_size[2], ct=chunk_size[0], ch=chunk_size[1], cw=chunk_size[2]) + + return x + + +# TEST +def generate_data(batch_size, seqlen, num_head, head_dim, dtype): + random.seed(0) + torch.manual_seed(0) + torch.cuda.manual_seed(0) + device = torch.cuda.current_device() + + q = torch.randn((batch_size, seqlen, num_head, head_dim), requires_grad=True).to(dtype=dtype, device='cuda') + k = torch.randn((batch_size, seqlen, num_head, head_dim), requires_grad=True).to(dtype=dtype, device='cuda') + v = torch.randn((batch_size, seqlen, num_head, head_dim), requires_grad=True).to(dtype=dtype, device='cuda') + print(f"q.shape: {q.shape}, k.shape: {k.shape}, v.shape: {v.shape}") + cu_seqlens = torch.arange(0, q.shape[0] * q.shape[1] + 1, q.shape[1], dtype=torch.int32, device='cuda') + max_seqlen = q.shape[1] + q = rearrange(q, "b s ... -> (b s) ...") + k = rearrange(k, "b s ... -> (b s) ...") + v = rearrange(v, "b s ... -> (b s) ...") + + return q, k, v, cu_seqlens, max_seqlen + + +def test_attn_varlen_moba_speed(batch, head, seqlen, head_dim, moba_chunk_size, moba_topk, dtype=torch.bfloat16, select_mode='threshold', simsum_threshold=0.25, threshold_type='query_head'): + """Speed test comparing flash_attn vs moba_attention""" + # Get data + q, k, v, cu_seqlen, max_seqlen = generate_data(batch, seqlen, head, head_dim, dtype) + print(f"batch:{batch} head:{head} seqlen:{seqlen} chunk:{moba_chunk_size} topk:{moba_topk} select_mode: {select_mode} simsum_threshold:{simsum_threshold}") + vo_grad = torch.randn_like(q) + + # Warmup + warmup_iters = 3 + perf_test_iters = 10 + + # Warmup + for _ in range(warmup_iters): + o = flash_attn_varlen_func(q, k, v, cu_seqlen, cu_seqlen, max_seqlen, max_seqlen, causal=False) + torch.autograd.backward(o, vo_grad) + + torch.cuda.synchronize() + start_flash = time.perf_counter() + for _ in range(perf_test_iters): + o = flash_attn_varlen_func(q, k, v, cu_seqlen, cu_seqlen, max_seqlen, max_seqlen, causal=False) + torch.autograd.backward(o, vo_grad) + + torch.cuda.synchronize() + time_flash = (time.perf_counter() - start_flash) / perf_test_iters * 1000 + + # Warmup + for _ in range(warmup_iters): + om = moba_attn_varlen(q, k, v, cu_seqlen, max_seqlen, moba_chunk_size=moba_chunk_size, moba_topk=moba_topk, select_mode=select_mode, simsum_threshold=simsum_threshold, threshold_type=threshold_type) + torch.autograd.backward(om, vo_grad) + + torch.cuda.synchronize() + start_moba = time.perf_counter() + for _ in range(perf_test_iters): + om = moba_attn_varlen(q, k, v, cu_seqlen, max_seqlen, moba_chunk_size=moba_chunk_size, moba_topk=moba_topk, select_mode=select_mode, simsum_threshold=simsum_threshold, threshold_type=threshold_type) + torch.autograd.backward(om, vo_grad) + + torch.cuda.synchronize() + time_moba = (time.perf_counter() - start_moba) / perf_test_iters * 1000 + + print(f"Flash: {time_flash:.2f}ms, MoBA: {time_moba:.2f}ms") + print(f"Speedup: {time_flash / time_moba:.2f}x") + + + +if __name__ == "__main__": + """ + CUDA_VISIBLE_DEVICES=1 \ + python -u csrc/attn/vmoba_attn/vmoba/vmoba.py + """ + test_attn_varlen_moba_speed(batch=1, head=12, seqlen=32760, head_dim=128, moba_chunk_size=32760 // 3 // 6 // 4, moba_topk=3, select_mode='threshold', simsum_threshold=0.3, threshold_type='query_head') diff --git a/csrc/fastvideo_kernel/tests/support_flex_sta.py b/csrc/fastvideo_kernel/tests/support_flex_sta.py new file mode 100644 index 000000000..d7d9528a6 --- /dev/null +++ b/csrc/fastvideo_kernel/tests/support_flex_sta.py @@ -0,0 +1,71 @@ +from typing import Tuple + +import torch +from torch import BoolTensor, IntTensor +from torch.nn.attention.flex_attention import create_block_mask + +# Peiyuan: This is neccesay. Dont know why. see https://github.com/pytorch/pytorch/issues/135028 +torch._inductor.config.realize_opcount_threshold = 100 + + +def generate_sta_mask(canvas_twh, kernel_twh, tile_twh, text_length): + """Generates a 3D NATTEN attention mask with a given kernel size. + + Args: + canvas_t: The time dimension of the canvas. + canvas_h: The height of the canvas. + canvas_w: The width of the canvas. + kernel_t: The time dimension of the kernel. + kernel_h: The height of the kernel. + kernel_w: The width of the kernel. + """ + canvas_t, canvas_h, canvas_w = canvas_twh + kernel_t, kernel_h, kernel_w = kernel_twh + tile_t_size, tile_h_size, tile_w_size = tile_twh + total_tile_size = tile_t_size * tile_h_size * tile_w_size + canvas_tile_t, canvas_tile_h, canvas_tile_w = canvas_t // tile_t_size, canvas_h // tile_h_size, canvas_w // tile_w_size + img_seq_len = canvas_t * canvas_h * canvas_w + + def get_tile_t_x_y(idx: IntTensor) -> Tuple[IntTensor, IntTensor, IntTensor]: + tile_id = idx // total_tile_size + tile_t = tile_id // (canvas_tile_h * canvas_tile_w) + tile_h = (tile_id % (canvas_tile_h * canvas_tile_w)) // canvas_tile_w + tile_w = tile_id % canvas_tile_w + return tile_t, tile_h, tile_w + + def sta_mask_3d( + b: IntTensor, + h: IntTensor, + q_idx: IntTensor, + kv_idx: IntTensor, + ) -> BoolTensor: + q_t_tile, q_x_tile, q_y_tile = get_tile_t_x_y(q_idx) + kv_t_tile, kv_x_tile, kv_y_tile = get_tile_t_x_y(kv_idx) + # kernel nominally attempts to center itself on the query, but kernel center + # is clamped to a fixed distance (kernel half-length) from the canvas edge + kernel_center_t = q_t_tile.clamp(kernel_t // 2, (canvas_tile_t - 1) - kernel_t // 2) + kernel_center_x = q_x_tile.clamp(kernel_h // 2, (canvas_tile_h - 1) - kernel_h // 2) + kernel_center_y = q_y_tile.clamp(kernel_w // 2, (canvas_tile_w - 1) - kernel_w // 2) + time_mask = (kernel_center_t - kv_t_tile).abs() <= kernel_t // 2 + hori_mask = (kernel_center_x - kv_x_tile).abs() <= kernel_h // 2 + vert_mask = (kernel_center_y - kv_y_tile).abs() <= kernel_w // 2 + image_mask = (q_idx < img_seq_len) & (kv_idx < img_seq_len) + image_to_text_mask = (q_idx < img_seq_len) & (kv_idx >= img_seq_len) & (kv_idx < img_seq_len + text_length) + text_to_all_mask = (q_idx >= img_seq_len) & (kv_idx < img_seq_len + text_length) + return (image_mask & time_mask & hori_mask & vert_mask) | image_to_text_mask | text_to_all_mask + + sta_mask_3d.__name__ = f"natten_3d_c{canvas_t}x{canvas_w}x{canvas_h}_k{kernel_t}x{kernel_w}x{kernel_h}" + return sta_mask_3d + + +def get_sliding_tile_attention_mask(kernel_size, tile_size, img_size, text_length, device, text_max_len=256): + img_seq_len = img_size[0] * img_size[1] * img_size[2] + image_mask = generate_sta_mask(img_size, kernel_size, tile_size, text_length) + mask = create_block_mask(image_mask, + B=None, + H=None, + Q_LEN=img_seq_len + text_max_len, + KV_LEN=img_seq_len + text_max_len, + device=device, + _compile=True) + return mask diff --git a/csrc/fastvideo_kernel/tests/test_correctness.py b/csrc/fastvideo_kernel/tests/test_correctness.py new file mode 100644 index 000000000..a2f29daa5 --- /dev/null +++ b/csrc/fastvideo_kernel/tests/test_correctness.py @@ -0,0 +1,63 @@ +import torch +import sys +import os +from tqdm import tqdm + +# Local support import +from .support_flex_sta import get_sliding_tile_attention_mask + +# USE OUR NEW PACKAGE! +from fastvideo_kernel import sliding_tile_attention +from torch.nn.attention.flex_attention import flex_attention + +flex_attention = torch.compile(flex_attention, dynamic=False) + +def flex_test(Q, K, V, kernel_size): + mask = get_sliding_tile_attention_mask(kernel_size, (6, 8, 8), (18, 48, 80), 0, 'cuda', 0) + output = flex_attention(Q, K, V, block_mask=mask) + return output + +def h100_fwd_kernel_test(Q, K, V, kernel_size): + # Using the same parameters as the original test + o = sliding_tile_attention(Q, K, V, [kernel_size] * 24, 0, False, '18x48x80') + return o + +def generate_tensor(shape, mean, std, dtype, device): + tensor = torch.randn(shape, dtype=dtype, device=device) + magnitude = torch.norm(tensor, dim=-1, keepdim=True) + scaled_tensor = tensor * (torch.randn(magnitude.shape, dtype=dtype, device=device) * std + mean) / magnitude + return scaled_tensor.contiguous() + +def check_correctness(b, h, n, d, causal, mean, std, num_iterations=2): + print(f"Running correctness check: batch={b}, heads={h}, seq_len={n}, dim={d}") + kernel_size_ls = [(3, 3, 5), (3, 1, 10)] + + for kernel_size in kernel_size_ls: + print(f"Testing kernel_size: {kernel_size}") + for xi in tqdm(range(num_iterations)): + torch.manual_seed(xi) + Q = generate_tensor((b, h, n, d), mean, std, torch.bfloat16, 'cuda') + K = generate_tensor((b, h, n, d), mean, std, torch.bfloat16, 'cuda') + V = generate_tensor((b, h, n, d), mean, std, torch.bfloat16, 'cuda') + + tk_o = h100_fwd_kernel_test(Q, K, V, kernel_size) + pt_o = flex_test(Q, K, V, kernel_size) + + diff = pt_o - tk_o + abs_diff = torch.abs(diff) + max_d = torch.max(abs_diff).item() + avg_d = torch.sum(abs_diff).item() / (b * h * n * d) + + if max_d > 0.1: + print(f"Warning: Large diff detected! max={max_d}, avg={avg_d}") + + print("\n✅ TEST COMPLETE: New package matches FlexAttention behavior.") + +if __name__ == "__main__": + b, h, d = 2, 24, 128 + n = 69120 + causal = False + mean = 1e-1 + std = 10 + + check_correctness(b, h, n, d, causal, mean, std, num_iterations=2) diff --git a/csrc/fastvideo_kernel/tests/test_vmoba_correctness.py b/csrc/fastvideo_kernel/tests/test_vmoba_correctness.py new file mode 100644 index 000000000..c5d1a7f72 --- /dev/null +++ b/csrc/fastvideo_kernel/tests/test_vmoba_correctness.py @@ -0,0 +1,97 @@ +# SPDX-License-Identifier: Apache-2.0 + +import torch +import pytest +import random +from fastvideo_kernel.vmoba import moba_attn_varlen + +def generate_test_data(batch_size, total_seqlen, num_heads, head_dim, dtype, device="cuda"): + """ + Generates random data for testing the variable-length attention function. + """ + torch.manual_seed(42) + random.seed(42) + torch.cuda.manual_seed_all(42) + + # Generate sequence lengths for each item in the batch + if batch_size > 1: + # Ensure sequence lengths are reasonably distributed + avg_seqlen = total_seqlen // batch_size + seqlens = [random.randint(avg_seqlen // 2, avg_seqlen + avg_seqlen // 2) for _ in range(batch_size - 1)] + remaining_len = total_seqlen - sum(seqlens) + if remaining_len > 0: + seqlens.append(remaining_len) + else: # Adjust if sum exceeds total_seqlen + seqlens.append(avg_seqlen) + current_sum = sum(seqlens) + seqlens[-1] -= (current_sum - total_seqlen) + # Ensure all lengths are positive + seqlens = [max(1, s) for s in seqlens] + # Final adjustment to match total_seqlen + seqlens[-1] += total_seqlen - sum(seqlens) + + else: + seqlens = [total_seqlen] + + cu_seqlens = torch.tensor([0] + list(torch.cumsum(torch.tensor(seqlens), 0)), device=device, dtype=torch.int32) + max_seqlen = max(seqlens) if seqlens else 0 + + q = torch.randn((total_seqlen, num_heads, head_dim), dtype=dtype, device=device, requires_grad=False) + k = torch.randn((total_seqlen, num_heads, head_dim), dtype=dtype, device=device, requires_grad=False) + v = torch.randn((total_seqlen, num_heads, head_dim), dtype=dtype, device=device, requires_grad=False) + + return q, k, v, cu_seqlens, max_seqlen + + +@pytest.mark.parametrize("batch_size", [1, 2]) +@pytest.mark.parametrize("total_seqlen", [512, 1024]) +@pytest.mark.parametrize("num_heads", [8]) +@pytest.mark.parametrize("head_dim", [64]) +@pytest.mark.parametrize("moba_chunk_size", [64]) +@pytest.mark.parametrize("moba_topk", [2, 4]) +@pytest.mark.parametrize("select_mode", ["topk", "threshold"]) +@pytest.mark.parametrize("threshold_type", ["query_head", "head_global", "overall"]) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +def test_moba_attn_varlen_forward( + batch_size, total_seqlen, num_heads, head_dim, moba_chunk_size, moba_topk, select_mode, threshold_type, dtype +): + """ + Tests the forward pass of moba_attn_varlen for basic correctness. + It checks output shape, dtype, and for the presence of NaNs/Infs. + """ + if dtype == torch.float32: + pytest.skip("float32 is not supported in flash attention") + + q, k, v, cu_seqlens, max_seqlen = generate_test_data( + batch_size, total_seqlen, num_heads, head_dim, dtype + ) + + # Ensure chunk size is not larger than the smallest sequence length + min_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).min().item() + if moba_chunk_size > min_seqlen: + pytest.skip("moba_chunk_size is larger than the minimum sequence length in the batch") + + try: + output = moba_attn_varlen( + q=q, + k=k, + v=v, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, + moba_chunk_size=moba_chunk_size, + moba_topk=moba_topk, + select_mode=select_mode, + threshold_type=threshold_type, + simsum_threshold=0.5, # A reasonable default for threshold mode + ) + except Exception as e: + pytest.fail(f"moba_attn_varlen forward pass failed with exception: {e}") + + # 1. Check output shape + assert output.shape == q.shape, f"Expected output shape {q.shape}, but got {output.shape}" + + # 2. Check output dtype + assert output.dtype == q.dtype, f"Expected output dtype {q.dtype}, but got {output.dtype}" + + # 3. Check for NaNs or Infs in the output + assert torch.all(torch.isfinite(output)), "Output contains NaN or Inf values" diff --git a/fastvideo/attention/backends/sliding_tile_attn.py b/fastvideo/attention/backends/sliding_tile_attn.py index 77ab57a8b..51dba3652 100644 --- a/fastvideo/attention/backends/sliding_tile_attn.py +++ b/fastvideo/attention/backends/sliding_tile_attn.py @@ -5,7 +5,7 @@ from typing import Any import torch from einops import rearrange -from st_attn import sliding_tile_attention +from fastvideo_kernel import sliding_tile_attention import fastvideo.envs as envs from fastvideo.attention.backends.abstract import (AttentionBackend, diff --git a/fastvideo/attention/backends/video_sparse_attn.py b/fastvideo/attention/backends/video_sparse_attn.py index 15c1c6ba0..0adc56724 100644 --- a/fastvideo/attention/backends/video_sparse_attn.py +++ b/fastvideo/attention/backends/video_sparse_attn.py @@ -6,7 +6,7 @@ from dataclasses import dataclass import torch try: - from vsa import video_sparse_attn + from fastvideo_kernel import video_sparse_attn except ImportError: video_sparse_attn = None diff --git a/fastvideo/attention/backends/vmoba.py b/fastvideo/attention/backends/vmoba.py index df55f7d72..1b27cd87a 100644 --- a/fastvideo/attention/backends/vmoba.py +++ b/fastvideo/attention/backends/vmoba.py @@ -6,8 +6,8 @@ from dataclasses import dataclass import torch from einops import rearrange -from csrc.attn.vmoba_attn.vmoba import (moba_attn_varlen, process_moba_input, - process_moba_output) +from fastvideo_kernel import (moba_attn_varlen, process_moba_input, + process_moba_output) from fastvideo.attention.backends.abstract import (AttentionBackend, AttentionImpl, AttentionMetadata, diff --git a/fastvideo/platforms/cuda.py b/fastvideo/platforms/cuda.py index c1f22df8a..c93be7f03 100644 --- a/fastvideo/platforms/cuda.py +++ b/fastvideo/platforms/cuda.py @@ -126,7 +126,7 @@ class CudaPlatformBase(Platform): logger.info("Selected backend: %s", selected_backend) if selected_backend == AttentionBackendEnum.SLIDING_TILE_ATTN: try: - from st_attn import sliding_tile_attention # noqa: F401 + from fastvideo_kernel import sliding_tile_attention # noqa: F401 from fastvideo.attention.backends.sliding_tile_attn import ( # noqa: F401 SlidingTileAttentionBackend) @@ -167,7 +167,7 @@ class CudaPlatformBase(Platform): ) elif selected_backend == AttentionBackendEnum.VIDEO_SPARSE_ATTN: try: - from vsa import block_sparse_attn # noqa: F401 + from fastvideo_kernel import video_sparse_attn # noqa: F401 from fastvideo.attention.backends.video_sparse_attn import ( # noqa: F401 VideoSparseAttentionBackend) @@ -186,8 +186,7 @@ class CudaPlatformBase(Platform): elif selected_backend == AttentionBackendEnum.VMOBA_ATTN: try: - from csrc.attn.vmoba_attn.vmoba import ( # noqa: F401 - moba_attn_varlen) + from fastvideo_kernel import moba_attn_varlen # noqa: F401 from fastvideo.attention.backends.vmoba import ( # noqa: F401 VMOBAAttentionBackend) logger.info("Using Video MOBA Attention backend.")