diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index dbbc1714a..8550c254e 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -129,11 +129,7 @@ steps: queue: "default" - path: - "fastvideo/**" - - "csrc/attn/video_sparse_attn/**" - - "csrc/attn/video_sparse_attn/tk/**" - - "csrc/attn/video_sparse_attn/setup.py" - - "csrc/attn/video_sparse_attn/config_vsa.py" - - "csrc/attn/video_sparse_attn/vsa.cpp" + - "fastvideo-kernel/**" - "pyproject.toml" - "docker/Dockerfile.python3.12" config: @@ -145,10 +141,7 @@ steps: queue: "default" - path: - "fastvideo/**" - - "csrc/attn/sliding_tile_attn/**" - - "csrc/attn/sliding_tile_attn/setup.py" - - "csrc/attn/sliding_tile_attn/config_sta.py" - - "csrc/attn/sliding_tile_attn/st_attn.cpp" + - "fastvideo-kernel/**" - "pyproject.toml" - "docker/Dockerfile.python3.12" config: @@ -159,10 +152,7 @@ steps: agents: queue: "default" - path: - - "csrc/attn/sliding_tile_attn/**" - - "csrc/attn/sliding_tile_attn/setup.py" - - "csrc/attn/sliding_tile_attn/config_sta.py" - - "csrc/attn/sliding_tile_attn/st_attn.cpp" + - "fastvideo-kernel/**" - "pyproject.toml" - "docker/Dockerfile.python3.12" config: @@ -173,12 +163,7 @@ steps: agents: queue: "default" - path: - - "csrc/attn/video_sparse_attn/**" - - "csrc/attn/video_sparse_attn/tk/**" - - "csrc/attn/tests/test_vsa.py" - - "csrc/attn/video_sparse_attn/setup.py" - - "csrc/attn/video_sparse_attn/config_vsa.py" - - "csrc/attn/video_sparse_attn/vsa.cpp" + - "fastvideo-kernel/**" - "pyproject.toml" - "docker/Dockerfile.python3.12" config: @@ -189,7 +174,7 @@ steps: agents: queue: "default" - path: - - "csrc/attn/vmoba_attn/**" + - "fastvideo-kernel/**" - "pyproject.toml" - "docker/Dockerfile.python3.12" config: @@ -200,7 +185,7 @@ steps: agents: queue: "default" - path: - - "csrc/attn/vmoba_attn/vmoba/**" + - "fastvideo-kernel/**" - "fastvideo/attention/backends/vmoba.py" - "pyproject.toml" - "docker/Dockerfile.python3.12" diff --git a/.github/workflows/fastvideo-kernel-publish.yml b/.github/workflows/fastvideo-kernel-publish.yml index 4a1e4fa44..973a900e8 100644 --- a/.github/workflows/fastvideo-kernel-publish.yml +++ b/.github/workflows/fastvideo-kernel-publish.yml @@ -5,7 +5,7 @@ on: branches: - main paths: - - "csrc/fastvideo_kernel/pyproject.toml" + - "fastvideo-kernel/pyproject.toml" workflow_dispatch: jobs: @@ -23,13 +23,15 @@ jobs: - name: Check if version changed id: check-version run: | - cd csrc/fastvideo_kernel + cd fastvideo-kernel # Get current commit's version from pyproject.toml - NEW_VERSION=$(grep -oP 'version\s*=\s*"\K[^"]+' pyproject.toml) + # Use ^ to match start of line to avoid matching minimum-version + 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") + # Note: git show expects path relative to repo root + OLD_VERSION=$(git show HEAD~1:fastvideo-kernel/pyproject.toml | grep -oP '^version\s*=\s*"\K[^"]+' || echo "0.0.0") echo "Old version: $OLD_VERSION" if [ "$NEW_VERSION" != "$OLD_VERSION" ]; then @@ -51,15 +53,18 @@ jobs: fail-fast: false matrix: os: [ubuntu-22.04] - python-version: ['3.10', '3.11', '3.12', '3.13'] + python-version: ['3.10', '3.11', '3.12'] 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' + # - 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' + - torch-version: '2.9.1' cuda-version: '12.8.0' torch-cuda-short: 'cu128' @@ -138,30 +143,30 @@ jobs: run: | export PYTHONPATH=$GITHUB_WORKSPACE:$PYTHONPATH - pip install setuptools ninja packaging wheel triton + pip install setuptools ninja packaging wheel triton scikit-build-core cmake build - cd csrc/fastvideo_kernel + cd 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 + export TORCH_CUDA_ARCH_LIST="9.0a" + + # Build standard wheel (no local version suffix) for PyPI + python -m build --wheel --outdir dist + + # Fix the wheel to be manylinux compliant + pip install auditwheel + # Target manylinux_2_35 (Ubuntu 22.04 native) + auditwheel repair dist/*.whl --plat manylinux_2_35_x86_64 -w fixed_dist + # Move fixed wheels back to dist for upload consistency + rm dist/*.whl + mv fixed_dist/*.whl dist/ - name: Upload wheel artifact + # Only upload if it's the "main" CUDA version we want on PyPI + # We upload all to artifacts for inspection/GH releases, but give them distinct artifact names uses: actions/upload-artifact@v4 with: - name: ${{ env.wheel_name }}-py${{ matrix.python-version }} - path: csrc/fastvideo_kernel/dist/*.whl + name: fastvideo_kernel-py${{ matrix.python-version }}-${{ matrix.torch-cuda.torch-cuda-short }}-torch${{ matrix.torch-cuda.torch-version }} + path: fastvideo-kernel/dist/*.whl retention-days: 90 publish_package: @@ -179,58 +184,22 @@ jobs: with: python-version: '3.10' - - name: Install CUDA 12.4.1 - uses: Jimver/cuda-toolkit@v0.2.21 - id: cuda-toolkit + - name: Download PyPI wheels + uses: actions/download-artifact@v4 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)" + path: fastvideo-kernel/dist/ + pattern: 'fastvideo_kernel-py*' + merge-multiple: true - name: Build source distribution run: | - export PYTHONPATH=$GITHUB_WORKSPACE:$PYTHONPATH + pip install build scikit-build-core cmake ninja - pip install setuptools ninja packaging wheel triton - - cd csrc/fastvideo_kernel - git submodule update --init --recursive - python setup.py sdist --dist-dir=dist + cd fastvideo-kernel + # We don't need full CUDA/Torch to just package the source (sdist) + python -m build --sdist --outdir dist - name: Publish release distributions to PyPI uses: pypa/gh-action-pypi-publish@release/v1 with: - packages-dir: csrc/fastvideo_kernel/dist/ + packages-dir: fastvideo-kernel/dist/ diff --git a/.gitmodules b/.gitmodules index bdb04686d..c51738164 100644 --- a/.gitmodules +++ b/.gitmodules @@ -1,7 +1,3 @@ -[submodule "csrc/attn/video_sparse_attn/tk"] - path = csrc/attn/video_sparse_attn/tk - url = https://github.com/HazyResearch/ThunderKittens.git - -[submodule "csrc/attn/sliding_tile_attn/tk"] - path = csrc/attn/sliding_tile_attn/tk +[submodule "fastvideo-kernel/include/tk"] + path = fastvideo-kernel/include/tk url = https://github.com/HazyResearch/ThunderKittens.git diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 1eeb53425..db9151d38 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -4,7 +4,7 @@ default_stages: exclude: | (?x)( fastvideo/third_party/.*| - csrc/.*| + fastvideo-kernel/.*| assets/.*| tests/.*| demo/.*| diff --git a/csrc/attn/README.md b/csrc/attn/README.md deleted file mode 100644 index 461e5bde5..000000000 --- a/csrc/attn/README.md +++ /dev/null @@ -1,113 +0,0 @@ - - -# Attention Kernel Used in FastVideo - - - -## Video Sparse Attention (VSA) - -### Installation -We support H100 (via TK) and any other GPU (via triton) for VSA. -```bash -git submodule update --init --recursive -python setup_vsa.py install -``` - - -If you encounter error during installation, try below: -Install C++20 for ThunderKittens: -```bash -sudo apt update -sudo apt install gcc-11 g++-11 - -sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11 - -sudo apt update -sudo apt install clang-11 -``` -(If you use CUDA12.8) -```bash -export CUDA_HOME=/usr/local/cuda-12.8 -export PATH=${CUDA_HOME}/bin:${PATH} -export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH -``` - -### Verify if you have successfully installed - -```bash -# test numerical -python tests/test_vsa.py -# (For H100) test speed -python benchmarks/bench_vsa_hopper.py -``` -bench_vsa_hopper.py should print something like this: -```bash -Using topk=76 kv blocks per q block (out of 768 total kv blocks) - -=== BLOCK SPARSE ATTENTION BENCHMARK === -Block Sparse Forward - TFLOPS: 5622.26 -Block Sparse Backward - TFLOPS: 3865.68 -``` - - -## Sliding Tile Attention (STA) -We only support H100 for STA. -```bash -git submodule update --init --recursive -python setup_sta.py install -``` - - - - -### Usage -End-2-end inference with FastVideo: -```bash -bash scripts/inference/v1_inference_wan_STA.sh -``` - -If you want to use sliding tile attention in your custom model: -```python -from st_attn import sliding_tile_attention -# assuming video size (T, H, W) = (30, 48, 80), text tokens = 256 with padding. -# q, k, v: [batch_size, num_heads, seq_length, head_dim], seq_length = T*H*W + 256 -# a tile is a cube of size (6, 8, 8) -# window_size in tiles: [(window_t, window_h, window_w), (..)...]. For example, window size (3, 3, 3) means a query can attend to (3x6, 3x8, 3x8) = (18, 24, 24) tokens out of the total 30x48x80 video. -# text_length: int ranging from 0 to 256 -# If your attention contains text token (Hunyuan) -out = sliding_tile_attention(q, k, v, window_size, text_length) -# If your attention does not contain text token (StepVideo) -out = sliding_tile_attention(q, k, v, window_size, 0, False) -``` - - -### Test -```bash -python tests/test_sta.py # test STA -python tests/test_vsa.py # test VSA -``` -### Benchmark -```bash -python benchmarks/bench_sta.py -``` - - -### How Does STA Work? -We give a demo for 2D STA with window size (6,6) operating on a (10, 10) image. - - -https://github.com/user-attachments/assets/f3b6dd79-7b43-4b60-a0fa-3d6495ec5747 - -## Why is STA Fast? -2D/3D Sliding Window Attention (SWA) creates many mixed blocks in the attention map. Even though mixed blocks have less output value,a mixed block is significantly slower than a dense block due to the GPU-unfriendly masking operation. - -STA removes mixed blocks. - - -
- -
- -## Acknowledgement - -We learned or reuse code from FlexAtteniton, NATEN, and ThunderKittens. diff --git a/csrc/attn/benchmarks/bench_sta.py b/csrc/attn/benchmarks/bench_sta.py deleted file mode 100644 index 8e0c8afc1..000000000 --- a/csrc/attn/benchmarks/bench_sta.py +++ /dev/null @@ -1,145 +0,0 @@ -import os -from collections import defaultdict - -import matplotlib.pyplot as plt -import numpy as np -import torch -from st_attn import sliding_tile_attention -from triton.testing import do_bench - - -def flops(batch, seqlen, nheads, headdim, causal, mode="fwd"): - assert mode in ["fwd", "bwd", "fwd_bwd"] - f = 4 * batch * seqlen**2 * nheads * headdim // (2 if causal else 1) - return f if mode == "fwd" else (2.5 * f if mode == "bwd" else 3.5 * f) - - -def compute_TFLOPS(flops, ms): - flops = flops / 1e12 - ms = ms / 1e3 - return flops / ms - - -def benchmark_attention(configurations): - results = {'fwd': defaultdict(list), 'bwd': defaultdict(list)} - - for B, H, N, D, causal, dit_seq_shape, window_size in configurations: - print("=" * 60) - print(f"Timing forward and backward pass for B={B}, H={H}, N={N}, D={D}, causal={causal}") - - q = torch.randn(B, H, N, D, dtype=torch.bfloat16, device='cuda', requires_grad=False).contiguous() - k = torch.randn(B, H, N, D, dtype=torch.bfloat16, device='cuda', requires_grad=False).contiguous() - v = torch.randn(B, H, N, D, dtype=torch.bfloat16, device='cuda', requires_grad=False).contiguous() - - # grad_output = torch.randn_like(q, requires_grad=False).contiguous() - # qg = torch.zeros_like(q, requires_grad=False, dtype=torch.float).contiguous() - # kg = torch.zeros_like(k, requires_grad=False, dtype=torch.float).contiguous() - # vg = torch.zeros_like(v, requires_grad=False, dtype=torch.float).contiguous() - - - # # Warmup for forward pass - # for _ in range(10): - # o = sliding_tile_attention(q, k, v, [[3, 6, 10]] * 24, 0, False, dit_seq_shape) - - # # Time the forward pass - # for i in range(10): - # start_events_fwd[i].record() - # o = sliding_tile_attention(q, k, v, [[3, 6, 10]] * 24, 0, False, dit_seq_shape) - # end_events_fwd[i].record() - ms = do_bench(lambda: sliding_tile_attention(q, k, v, [window_size] * 24, 0, False, dit_seq_shape)) - - # times_fwd = [s.elapsed_time(e) for s, e in zip(start_events_fwd, end_events_fwd)] - # time_us_fwd = np.mean(times_fwd) * 1000 - - tflops_fwd = compute_TFLOPS(flops(B, N, H, D, causal, 'fwd'), ms) - results['fwd'][(D, causal)].append((N, tflops_fwd)) - - print(f"Average time for forward pass (ms): {ms:.2f}") - print(f"Average TFLOPS: {tflops_fwd}") - print("-" * 60) - - # torch.cuda.empty_cache() - # torch.cuda.synchronize() - - # # Prepare for timing backward pass - # start_events_bwd = [torch.cuda.Event(enable_timing=True) for _ in range(10)] - # end_events_bwd = [torch.cuda.Event(enable_timing=True) for _ in range(10)] - - # # Warmup for backward pass - # for _ in range(10): - # qg, kg, vg = tk.mha_backward(q, k, v, o, l_vec, grad_output, causal) - - # # Time the backward pass - # for i in range(10): - # start_events_bwd[i].record() - # qg, kg, vg = tk.mha_backward(q, k, v, o, l_vec, grad_output, causal) - # end_events_bwd[i].record() - - # torch.cuda.synchronize() - # times_bwd = [s.elapsed_time(e) for s, e in zip(start_events_bwd, end_events_bwd)] - # time_us_bwd = np.mean(times_bwd) * 1000 - - # tflops_bwd = compute_TFLOPS(flops(B, N, H, D, causal, 'bwd'), ms) - # results['bwd'][(D, causal)].append((N, tflops_bwd)) - - # print(f"Average time for backward pass(ms): {ms:.2f}") - # print(f"Average TFLOPS: {tflops_bwd}") - # print("=" * 60) - - return results - - -def plot_results(results): - os.makedirs('benchmark_results', exist_ok=True) - for mode in ['fwd', 'bwd']: - for (D, causal), values in results[mode].items(): - seq_lens = [x[0] for x in values] - tflops = [x[1] for x in values] - - plt.figure(figsize=(10, 6)) - bars = plt.bar(range(len(seq_lens)), tflops, tick_label=seq_lens) - plt.xlabel('Sequence Length') - plt.ylabel('TFLOPS') - plt.title(f'{mode.upper()} Pass - Head Dim: {D}, Causal: {causal}') - plt.grid(True) - - # Adding the numerical y value on top of each bar - for bar in bars: - yval = bar.get_height() - plt.text(bar.get_x() + bar.get_width() / 2, yval, round(yval, 2), ha='center', va='bottom') - - filename = f'benchmark_results/{mode}_D{D}_causal{causal}.png' - plt.savefig(filename) - plt.close() - - -# Example list of configurations to test -configurations = [ - (2, 24, 69120, 128, False, '18x48x80', [3, 6, 10]), - (2, 24, 69120, 128, True, '18x48x80', [3, 6, 10]), - (2, 24, 82944, 128, False, '36x48x48', [3, 3, 6]), # Stepvideo - (2, 24, 82944, 128, True, '36x48x48', [3, 3, 6]), - # (16, 16, 768*16, 128, False), - # (16, 16, 768*2, 128, False), - # (16, 16, 768*4, 128, False), - # (16, 16, 768*8, 128, False), - # (16, 16, 768*16, 128, False), - # (16, 16, 768, 128, True), - # (16, 16, 768*2, 128, True), - # (16, 16, 768*4, 128, True), - # (16, 16, 768*8, 128, True), - # (16, 16, 768*16, 128, True), - # (16, 32, 768, 64, False), - # (16, 32, 768*2, 64, False), - # (16, 32, 768*4, 64, False), - # (16, 32, 768*8, 64, False), - # (16, 32, 768*16, 64, False), - # (16, 32, 768, 64, True), - # (16, 32, 768*2, 64, True), - # (16, 32, 768*4, 64, True), - # (16, 32, 768*8, 64, True), - # (16, 32, 768*16, 64, True), -] - -results = benchmark_attention(configurations) -# plot_results(results) diff --git a/csrc/attn/benchmarks/bench_vsa_hopper.py b/csrc/attn/benchmarks/bench_vsa_hopper.py deleted file mode 100644 index 893067376..000000000 --- a/csrc/attn/benchmarks/bench_vsa_hopper.py +++ /dev/null @@ -1,224 +0,0 @@ -import torch -import argparse -from triton.testing import do_bench -from vsa import block_sparse_fwd, block_sparse_bwd -from vsa import BLOCK_M, BLOCK_N -import triton -import numpy as np -import random - -def set_seed(seed: int = 42): - # Python random module - random.seed(seed) - - # NumPy - np.random.seed(seed) - - # PyTorch - torch.manual_seed(seed) - torch.cuda.manual_seed(seed) - torch.cuda.manual_seed_all(seed) # if using multi-GPU - -def parse_arguments(): - parser = argparse.ArgumentParser(description='Benchmark Block Sparse Attention') - parser.add_argument('--batch_size', type=int, default=1, help='Batch size') - parser.add_argument('--num_heads', type=int, default=12, help='Number of heads') - parser.add_argument('--head_dim', type=int, default=128, help='Head dimension') - parser.add_argument('--topk', type=int, default=None, help='Number of kv blocks each q block attends to') - parser.add_argument('--seq_lengths', type=int, nargs='+', default=[49152], help='Sequence lengths to benchmark') - return parser.parse_args() - -def create_input_tensors(batch, head, seq_len, headdim): - """Create random input tensors for attention.""" - q = torch.randn(batch, head, seq_len, headdim, dtype=torch.bfloat16, device="cuda") - k = torch.randn(batch, head, seq_len, headdim, dtype=torch.bfloat16, device="cuda") - v = torch.randn(batch, head, seq_len, headdim, dtype=torch.bfloat16, device="cuda") - return q, k, v - -def generate_block_sparse_pattern(bs, h, num_q_blocks, num_kv_blocks, k, device="cuda"): - """ - Generate a block sparse pattern where each q block attends to exactly k kv blocks. - - Args: - bs: batch size - h: number of heads - num_q_blocks: number of query blocks - num_kv_blocks: number of key-value blocks - k: number of kv blocks each q block attends to - device: device to create tensors on - - Returns: - q2k_block_sparse_index: [bs, h, num_q_blocks, k] - Contains the indices of kv blocks that each q block attends to. - q2k_block_sparse_num: [bs, h, num_q_blocks] - Contains the number of kv blocks that each q block attends to (all equal to k). - k2q_block_sparse_index: [bs, h, num_kv_blocks, num_q_blocks] - Contains the indices of q blocks that attend to each kv block. - k2q_block_sparse_num: [bs, h, num_kv_blocks] - Contains the number of q blocks that attend to each kv block. - block_sparse_mask: [bs, h, num_q_blocks, num_kv_blocks] - Binary mask where 1 indicates attention connection. - """ - # Ensure k is not larger than num_kv_blocks - k = min(k, num_kv_blocks) - - # Create random scores for sampling - scores = torch.rand(bs, h, num_q_blocks, num_kv_blocks, device=device) - - # Get top-k indices for each q block - _, q2k_block_sparse_index = torch.topk(scores, k, dim=-1) - q2k_block_sparse_index = q2k_block_sparse_index.to(torch.int32) - - # sort q2k_block_sparse_index - q2k_block_sparse_index, _ = torch.sort(q2k_block_sparse_index, dim=-1) - - # All q blocks attend to exactly k kv blocks - q2k_block_sparse_num = torch.full((bs, h, num_q_blocks), k, dtype=torch.int32, device=device) - - # Create the corresponding mask - block_sparse_mask = torch.zeros(bs, h, num_q_blocks, num_kv_blocks, dtype=torch.bool, device=device) - - # Fill in the mask based on the indices - for b in range(bs): - for head in range(h): - for q_idx in range(num_q_blocks): - kv_indices = q2k_block_sparse_index[b, head, q_idx] - block_sparse_mask[b, head, q_idx, kv_indices] = True - - # Create the reverse mapping (k2q) - # First, initialize lists to collect q indices for each kv block - k2q_indices_list = [[[] for _ in range(num_kv_blocks)] for _ in range(bs * h)] - - # Populate the lists based on q2k mapping - for b in range(bs): - for head in range(h): - flat_idx = b * h + head - for q_idx in range(num_q_blocks): - kv_indices = q2k_block_sparse_index[b, head, q_idx].tolist() - for kv_idx in kv_indices: - k2q_indices_list[flat_idx][kv_idx].append(q_idx) - - # Find the maximum number of q blocks that attend to any kv block - max_q_per_kv = 0 - for flat_idx in range(bs * h): - for kv_idx in range(num_kv_blocks): - max_q_per_kv = max(max_q_per_kv, len(k2q_indices_list[flat_idx][kv_idx])) - - # Create tensors for k2q mapping - k2q_block_sparse_index = torch.full((bs, h, num_kv_blocks, max_q_per_kv), -1, - dtype=torch.int32, device=device) - k2q_block_sparse_num = torch.zeros((bs, h, num_kv_blocks), - dtype=torch.int32, device=device) - - # Fill the tensors - for b in range(bs): - for head in range(h): - flat_idx = b * h + head - for kv_idx in range(num_kv_blocks): - q_indices = k2q_indices_list[flat_idx][kv_idx] - num_q = len(q_indices) - k2q_block_sparse_num[b, head, kv_idx] = num_q - if num_q > 0: - k2q_block_sparse_index[b, head, kv_idx, :num_q] = torch.tensor( - q_indices, dtype=torch.int32, device=device) - - return q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num, block_sparse_mask - -def benchmark_block_sparse_attention(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num, flops): - """Benchmark block sparse attention forward and backward passes.""" - print("\n=== BLOCK SPARSE ATTENTION BENCHMARK ===") - - # Forward pass - # Warm-up run - variable_block_sizes = torch.ones(q2k_block_sparse_index.shape[2], device=q.device).int() * BLOCK_M - o, l_vec = block_sparse_fwd(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, variable_block_sizes) - torch.cuda.synchronize() - - # Benchmark forward - fwd_time = do_bench( - lambda: block_sparse_fwd(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, variable_block_sizes), - warmup=5, - rep=20, - quantiles=None - ) - - sparse_tflops = flops / fwd_time * 1e-12 * 1e3 - print(f"Block Sparse Forward - TFLOPS: {sparse_tflops:.2f}") - - # Backward pass - grad_output = torch.randn_like(o) - - # Warm-up runs - for _ in range(5): - block_sparse_bwd(q, k, v, o, l_vec, grad_output, k2q_block_sparse_index, k2q_block_sparse_num, variable_block_sizes) - torch.cuda.synchronize() - - # Benchmark backward - bwd_time = do_bench( - lambda: block_sparse_bwd(q, k, v, o, l_vec, grad_output, k2q_block_sparse_index, k2q_block_sparse_num, variable_block_sizes), - warmup=5, - rep=20, - quantiles=None - ) - bwd_flops = 2.5 * flops # Approximation - - sparse_bwd_tflops = bwd_flops / bwd_time * 1e-12 * 1e3 - print(f"Block Sparse Backward - TFLOPS: {sparse_bwd_tflops:.2f}") - - return sparse_tflops, sparse_bwd_tflops - -def main(): - args = parse_arguments() - - set_seed(42) - - # Extract parameters - batch = args.batch_size - head = args.num_heads - headdim = args.head_dim - - print(f"Block Sparse Attention Benchmark") - print(f"batch: {batch}, head: {head}, headdim: {headdim}") - - # Test with different sequence lengths - for seq_len in args.seq_lengths: - # Skip very long sequences if they might cause OOM - if seq_len > 16384 and batch > 1: - continue - - print("="*100) - print(f"\nSequence length: {seq_len}") - - # Calculate theoretical FLOPs for attention - flops = 4 * batch * head * headdim * seq_len * seq_len - - # Create input tensors - q, k, v = create_input_tensors(batch, head, seq_len, headdim) - - # Setup block sparse parameters - num_q_blocks = seq_len // BLOCK_M - num_kv_blocks = seq_len // BLOCK_N - - # Determine k value (number of kv blocks per q block) - topk = args.topk - if topk is None: - topk = num_kv_blocks // 10 # Default to ~90% sparsity if k is not specified - topk = max(1, topk) - print(f"Using topk={topk} kv blocks per q block (out of {num_kv_blocks} total kv blocks)") - - # Generate block sparse pattern - q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num, _ = generate_block_sparse_pattern( - batch, head, num_q_blocks, num_kv_blocks, topk, device="cuda") - - # Benchmark block sparse attention - sparse_fwd, sparse_bwd = benchmark_block_sparse_attention( - q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num, flops - ) - - # Print results - print("\n=== PERFORMANCE RESULTS ===") - print(f"Block Sparse Forward - TFLOPS: {sparse_fwd:.2f}") - print(f"Block Sparse Backward - TFLOPS: {sparse_bwd:.2f}") - -if __name__ == "__main__": - main() \ No newline at end of file diff --git a/csrc/attn/benchmarks/bench_vsa_triton.py b/csrc/attn/benchmarks/bench_vsa_triton.py deleted file mode 100644 index 125b2677c..000000000 --- a/csrc/attn/benchmarks/bench_vsa_triton.py +++ /dev/null @@ -1,217 +0,0 @@ -import torch -import argparse -import triton.testing -from vsa import block_sparse_attn -from vsa import BLOCK_M, BLOCK_N - -import numpy as np -import random - -def set_seed(seed: int = 42): - # Python random module - random.seed(seed) - - # NumPy - np.random.seed(seed) - - # PyTorch - torch.manual_seed(seed) - torch.cuda.manual_seed(seed) - torch.cuda.manual_seed_all(seed) # if using multi-GPU - -def parse_arguments(): - parser = argparse.ArgumentParser(description='Benchmark Block Sparse Attention') - parser.add_argument('--batch_size', type=int, default=1, help='Batch size') - parser.add_argument('--num_heads', type=int, default=12, help='Number of heads') - parser.add_argument('--head_dim', type=int, default=64, help='Head dimension') - parser.add_argument('--topk', type=int, default=None, help='Number of kv blocks each q block attends to') - parser.add_argument('--seq_lengths', type=int, nargs='+', default=[49152], help='Sequence lengths to benchmark') - return parser.parse_args() - -def create_input_tensors(batch, head, seq_len, headdim): - """Create random input tensors for attention.""" - q = torch.randn(batch, head, seq_len, headdim, dtype=torch.bfloat16, device="cuda") - k = torch.randn(batch, head, seq_len, headdim, dtype=torch.bfloat16, device="cuda") - v = torch.randn(batch, head, seq_len, headdim, dtype=torch.bfloat16, device="cuda") - return q, k, v - -def generate_block_sparse_pattern(bs, h, num_q_blocks, num_kv_blocks, k, device="cuda"): - """ - Generate a block sparse pattern where each q block attends to exactly k kv blocks. - - Args: - bs: batch size - h: number of heads - num_q_blocks: number of query blocks - num_kv_blocks: number of key-value blocks - k: number of kv blocks each q block attends to - device: device to create tensors on - - Returns: - q2k_block_sparse_index: [bs, h, num_q_blocks, k] - Contains the indices of kv blocks that each q block attends to. - q2k_block_sparse_num: [bs, h, num_q_blocks] - Contains the number of kv blocks that each q block attends to (all equal to k). - k2q_block_sparse_index: [bs, h, num_kv_blocks, num_q_blocks] - Contains the indices of q blocks that attend to each kv block. - k2q_block_sparse_num: [bs, h, num_kv_blocks] - Contains the number of q blocks that attend to each kv block. - block_sparse_mask: [bs, h, num_q_blocks, num_kv_blocks] - Binary mask where 1 indicates attention connection. - """ - # Ensure k is not larger than num_kv_blocks - k = min(k, num_kv_blocks) - - # Create random scores for sampling - scores = torch.rand(bs, h, num_q_blocks, num_kv_blocks, device=device) - - # Get top-k indices for each q block - _, q2k_block_sparse_index = torch.topk(scores, k, dim=-1) - q2k_block_sparse_index = q2k_block_sparse_index.to(torch.int32) - - # sort q2k_block_sparse_index - q2k_block_sparse_index, _ = torch.sort(q2k_block_sparse_index, dim=-1) - - # All q blocks attend to exactly k kv blocks - q2k_block_sparse_num = torch.full((bs, h, num_q_blocks), k, dtype=torch.int32, device=device) - - # Create the corresponding mask - block_sparse_mask = torch.zeros(bs, h, num_q_blocks, num_kv_blocks, dtype=torch.bool, device=device) - - # Fill in the mask based on the indices - for b in range(bs): - for head in range(h): - for q_idx in range(num_q_blocks): - kv_indices = q2k_block_sparse_index[b, head, q_idx] - block_sparse_mask[b, head, q_idx, kv_indices] = True - - # Create the reverse mapping (k2q) - # First, initialize lists to collect q indices for each kv block - k2q_indices_list = [[[] for _ in range(num_kv_blocks)] for _ in range(bs * h)] - - # Populate the lists based on q2k mapping - for b in range(bs): - for head in range(h): - flat_idx = b * h + head - for q_idx in range(num_q_blocks): - kv_indices = q2k_block_sparse_index[b, head, q_idx].tolist() - for kv_idx in kv_indices: - k2q_indices_list[flat_idx][kv_idx].append(q_idx) - - # Find the maximum number of q blocks that attend to any kv block - max_q_per_kv = 0 - for flat_idx in range(bs * h): - for kv_idx in range(num_kv_blocks): - max_q_per_kv = max(max_q_per_kv, len(k2q_indices_list[flat_idx][kv_idx])) - - # Create tensors for k2q mapping - k2q_block_sparse_index = torch.full((bs, h, num_kv_blocks, max_q_per_kv), -1, - dtype=torch.int32, device=device) - k2q_block_sparse_num = torch.zeros((bs, h, num_kv_blocks), - dtype=torch.int32, device=device) - - # Fill the tensors - for b in range(bs): - for head in range(h): - flat_idx = b * h + head - for kv_idx in range(num_kv_blocks): - q_indices = k2q_indices_list[flat_idx][kv_idx] - num_q = len(q_indices) - k2q_block_sparse_num[b, head, kv_idx] = num_q - if num_q > 0: - k2q_block_sparse_index[b, head, kv_idx, :num_q] = torch.tensor( - q_indices, dtype=torch.int32, device=device) - - return q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num, block_sparse_mask - -def benchmark_block_sparse_attention(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num, flops): - """Benchmark block sparse attention forward+backward pass.""" - print("\n=== BLOCK SPARSE ATTENTION FORWARD+BACKWARD BENCHMARK ===") - - # Combined forward+backward pass - # Warm-up run - q_fwd = q.clone().requires_grad_(True) - k_fwd = k.clone().requires_grad_(True) - v_fwd = v.clone().requires_grad_(True) - o = block_sparse_attn(q_fwd, k_fwd, v_fwd, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num) - grad_output = torch.randn_like(o) - o.backward(grad_output) - torch.cuda.synchronize() - - # Benchmark forward+backward - def forward_backward_fn(): - q_fwd = q.clone().requires_grad_(True) - k_fwd = k.clone().requires_grad_(True) - v_fwd = v.clone().requires_grad_(True) - o = block_sparse_attn(q_fwd, k_fwd, v_fwd, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num) - grad_output = torch.randn_like(o) - o.backward(grad_output) - - total_time = triton.testing.do_bench( - forward_backward_fn, - warmup=25, - rep=100, - return_mode='mean' - ) - - # Total flops for forward + backward (forward + 2.5x backward approximation) - total_flops = flops + 2.5 * flops # 3.5x the forward flops - sparse_tflops = total_flops / total_time * 1e-12 * 1e3 - print(f"Block Sparse Forward+Backward - TFLOPS: {sparse_tflops:.2f}") - - return sparse_tflops - -def main(): - args = parse_arguments() - - set_seed(42) - - # Extract parameters - batch = args.batch_size - head = args.num_heads - headdim = args.head_dim - - print(f"Block Sparse Attention Benchmark") - print(f"batch: {batch}, head: {head}, headdim: {headdim}") - - # Test with different sequence lengths - for seq_len in args.seq_lengths: - # Skip very long sequences if they might cause OOM - if seq_len > 16384 and batch > 1: - continue - - print("="*100) - print(f"\nSequence length: {seq_len}") - - # Calculate theoretical FLOPs for attention - flops = 4 * batch * head * headdim * seq_len * seq_len - - # Create input tensors - q, k, v = create_input_tensors(batch, head, seq_len, headdim) - - # Setup block sparse parameters - num_q_blocks = seq_len // BLOCK_M - num_kv_blocks = seq_len // BLOCK_N - - # Determine k value (number of kv blocks per q block) - topk = args.topk - if topk is None: - topk = num_kv_blocks // 10 # Default to ~90% sparsity if k is not specified - topk = max(1, topk) - print(f"Using topk={topk} kv blocks per q block (out of {num_kv_blocks} total kv blocks)") - - # Generate block sparse pattern - q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num, _ = generate_block_sparse_pattern( - batch, head, num_q_blocks, num_kv_blocks, topk, device="cuda") - - # Benchmark block sparse attention - sparse_fwd = benchmark_block_sparse_attention( - q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num, flops - ) - - # Print results - print("\n=== PERFORMANCE RESULTS ===") - print(f"Block Sparse Forward+Backward - TFLOPS: {sparse_fwd:.2f}") - -if __name__ == "__main__": - main() \ No newline at end of file diff --git a/csrc/attn/sliding_tile_attn/MANIFEST.in b/csrc/attn/sliding_tile_attn/MANIFEST.in deleted file mode 100644 index cc38220c0..000000000 --- a/csrc/attn/sliding_tile_attn/MANIFEST.in +++ /dev/null @@ -1,2 +0,0 @@ -recursive-include tk * -include config_sta.py diff --git a/csrc/attn/sliding_tile_attn/README.md b/csrc/attn/sliding_tile_attn/README.md deleted file mode 100644 index b543d35b7..000000000 --- a/csrc/attn/sliding_tile_attn/README.md +++ /dev/null @@ -1,103 +0,0 @@ - -# Attention Kernel Used in FastVideo - -## Sliding Tile Attention (STA) -We support H100 (via TK) and any other GPU (via triton) for STA. - -### Installation -```bash -pip install st_attn -``` - -Install from source: - -```bash -git submodule update --init --recursive -python setup.py install -``` - -If you want to skip the compilation of the TK kernel and only use the Triton version, try below: - -```bash -SKIP_SM90_EXT=1 python setup.py install -or -SKIP_SM90_EXT=1 pip install --no-build-isolation . -``` - -If you encounter error during installation, try below: -Install C++20 for ThunderKittens: -```bash -sudo apt update -sudo apt install gcc-11 g++-11 - -sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11 - -sudo apt update -sudo apt install clang-11 -``` -(If you use CUDA12.8) -```bash -export CUDA_HOME=/usr/local/cuda-12.8 -export PATH=${CUDA_HOME}/bin:${PATH} -export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH -``` - -### Usage -End-2-end inference with FastVideo: -```bash -bash scripts/inference/v1_inference_wan_STA.sh -``` - -If you want to use sliding tile attention in your custom model: -```python -from st_attn import sliding_tile_attention -# assuming video size (T, H, W) = (30, 48, 80), text tokens = 256 with padding. -# q, k, v: [batch_size, num_heads, seq_length, head_dim], seq_length = T*H*W + 256 -# a tile is a cube of size (6, 8, 8) -# window_size in tiles: [(window_t, window_h, window_w), (..)...]. For example, window size (3, 3, 3) means a query can attend to (3x6, 3x8, 3x8) = (18, 24, 24) tokens out of the total 30x48x80 video. -# text_length: int ranging from 0 to 256 -# If your attention contains text token (Hunyuan) -out = sliding_tile_attention(q, k, v, window_size, text_length) -# If your attention does not contain text token (StepVideo) -out = sliding_tile_attention(q, k, v, window_size, 0, False) -``` - - -### Test -```bash -python ../tests/test_sta.py # test STA -``` -### Benchmark -```bash -python ../benchmarks/bench_sta.py -``` - - -### How Does STA Work? -We give a demo for 2D STA with window size (6,6) operating on a (10, 10) image. - - -https://github.com/user-attachments/assets/f3b6dd79-7b43-4b60-a0fa-3d6495ec5747 - - -## STA Configuration Logic -Here is a diagram of how the window is configured and passed through the FastVideo pipeline: - -
- -
- - -## Why is STA Fast? -2D/3D Sliding Window Attention (SWA) creates many mixed blocks in the attention map. Even though mixed blocks have less output value,a mixed block is significantly slower than a dense block due to the GPU-unfriendly masking operation. - -STA removes mixed blocks. - - -
- -
- -## Acknowledgement - -We learned or reuse code from FlexAtteniton, NATEN, and ThunderKittens. diff --git a/csrc/attn/sliding_tile_attn/config_sta.py b/csrc/attn/sliding_tile_attn/config_sta.py deleted file mode 100644 index f3008bcd3..000000000 --- a/csrc/attn/sliding_tile_attn/config_sta.py +++ /dev/null @@ -1,15 +0,0 @@ -### ADD TO THIS TO REGISTER NEW KERNELS -sources = { - 'st_attn': { - 'source_files': { - 'h100': 'st_attn/st_attn_h100.cu' # define these source files for each GPU target desired. - } - } -} - -### WHICH KERNELS DO WE WANT TO BUILD? -# (oftentimes during development work you don't need to redefine them all.) -kernels = ['st_attn'] - -### WHICH GPU TARGET DO WE WANT TO BUILD FOR? -target = 'h100' diff --git a/csrc/attn/sliding_tile_attn/setup.py b/csrc/attn/sliding_tile_attn/setup.py deleted file mode 100644 index fc55f38d7..000000000 --- a/csrc/attn/sliding_tile_attn/setup.py +++ /dev/null @@ -1,83 +0,0 @@ -import os -import subprocess - -from config_sta import kernels, sources, target -from setuptools import find_packages, setup -from torch.utils.cpp_extension import BuildExtension, CUDAExtension - -target = target.lower() - -# Package metadata -PACKAGE_NAME = "st_attn" -VERSION = "0.0.6" -AUTHOR = "Hao AI Lab" -DESCRIPTION = "Sliding Tile Atteniton Kernel Used in FastVideo" -URL = "https://github.com/hao-ai-lab/FastVideo/tree/main/csrc/sliding_tile_attention" - -# Set environment variables -tk_root = os.getenv('THUNDERKITTENS_ROOT', os.path.abspath(os.path.join(os.getcwd(), 'tk/'))) -python_include = subprocess.check_output(['python', '-c', - "import sysconfig; print(sysconfig.get_path('include'))"]).decode().strip() -torch_include = 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() -print('st_attn root:', tk_root) -print('Python include:', python_include) -print('Torch include directories:', torch_include) - -# CUDA flags -cuda_flags = [ - '-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' -] + torch_include.split() -cpp_flags = ['-std=c++20', '-O3'] - -if target == 'h100': - cuda_flags.append('-DKITTENS_HOPPER') - cuda_flags.append('-arch=sm_90a') -else: - raise ValueError(f'Target {target} not supported') - -source_files = ['st_attn.cpp'] -for k in kernels: - if target not in sources[k]['source_files']: - raise KeyError(f'Target {target} not found in source files for kernel {k}') - if isinstance(sources[k]['source_files'][target], list): - source_files.extend(sources[k]['source_files'][target]) - else: - source_files.append(sources[k]['source_files'][target]) - cpp_flags.append(f'-DTK_COMPILE_{k.replace(" ", "_").upper()}') - -ext_modules = [] - -if os.environ.get("SKIP_SM90_EXT", "0") != "1": - ext_modules.append( - CUDAExtension('st_attn_cuda', - sources=source_files, - extra_compile_args={ - 'cxx': cpp_flags, - 'nvcc': cuda_flags - }, - libraries=['cuda']) - ) -else: - print("ENV SKIP_SM90_EXT=1, skip st_attn_cuda compile") - -setup(name=PACKAGE_NAME, - version=VERSION, - author=AUTHOR, - description=DESCRIPTION, - url=URL, - packages=find_packages(), - ext_modules=ext_modules, - cmdclass={'build_ext': BuildExtension}, - classifiers=[ - "Programming Language :: Python :: 3", - "Environment :: GPU :: NVIDIA CUDA :: 12", - "License :: OSI Approved :: Apache Software License", - ], - python_requires='>=3.10', - install_requires=["torch>=2.5.0"]) diff --git a/csrc/attn/sliding_tile_attn/st_attn.cpp b/csrc/attn/sliding_tile_attn/st_attn.cpp deleted file mode 100644 index 41a63c17f..000000000 --- a/csrc/attn/sliding_tile_attn/st_attn.cpp +++ /dev/null @@ -1,23 +0,0 @@ -#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/attn/sliding_tile_attn/st_attn/__init__.py b/csrc/attn/sliding_tile_attn/st_attn/__init__.py deleted file mode 100644 index 89fb94bf8..000000000 --- a/csrc/attn/sliding_tile_attn/st_attn/__init__.py +++ /dev/null @@ -1,63 +0,0 @@ -import math - -import torch -from torch.utils.checkpoint import detach_variable -try: - from st_attn_cuda import sta_fwd -except ImportError: - sta_fwd = None - -try: - from st_attn.st_attn_triton import sliding_tile_attention_triton -except ImportError: - sliding_tile_attention_triton = None - -def sliding_tile_attention_SM90(q_all, k_all, v_all, window_size, text_length, has_text=True, dit_seq_shape='30x48x80'): - seq_length = q_all.shape[2] - dit_seq_shape_mapping = { - '30x48x80':1, - '36x48x48':2, - '18x48x80':3, - } - if has_text: - assert q_all.shape[ - 2] >= 115200 and q_all.shape[2] <= 115456, f"Unsupported {dit_seq_shape}, current shape is {q_all.shape}, only support '30x48x80' for HunyuanVideo" - assert q_all.shape[1] == len(window_size), "Number of heads must match the number of window sizes" - target_size = math.ceil(seq_length / 384) * 384 - pad_size = target_size - seq_length - if pad_size > 0: - q_all = torch.cat([q_all, q_all[:, :, -pad_size:]], dim=2) - k_all = torch.cat([k_all, k_all[:, :, -pad_size:]], dim=2) - v_all = torch.cat([v_all, v_all[:, :, -pad_size:]], dim=2) - else: - if dit_seq_shape == '36x48x48': # Stepvideo 204x768x68 - assert q_all.shape[2] == 82944 - elif dit_seq_shape == '18x48x80': # Wan 69x768x1280 - assert q_all.shape[2] == 69120 - else: - raise ValueError(f"Unsupported {dit_seq_shape}, current shape is {q_all.shape}, only support '36x48x48' for Stepvideo and '18x48x80' for Wan") - - kernel_aspect_ratio_flag = dit_seq_shape_mapping[dit_seq_shape] - hidden_states = torch.empty_like(q_all) - # This for loop is ugly. but it is actually quite efficient. The sequence dimension alone can already oversubscribe SMs - for head_index, (t_kernel, h_kernel, w_kernel) in enumerate(window_size): - for batch in range(q_all.shape[0]): - q_head, k_head, v_head, o_head = (q_all[batch:batch + 1, head_index:head_index + 1], - k_all[batch:batch + 1, - head_index:head_index + 1], v_all[batch:batch + 1, - head_index:head_index + 1], - hidden_states[batch:batch + 1, head_index:head_index + 1]) - - _ = sta_fwd(q_head, k_head, v_head, o_head, t_kernel, h_kernel, w_kernel, text_length, False, has_text, kernel_aspect_ratio_flag) - if has_text: - _ = sta_fwd(q_all, k_all, v_all, hidden_states, 3, 3, 3, text_length, True, True, kernel_aspect_ratio_flag) - return hidden_states[:, :, :seq_length] - -def sliding_tile_attention(q_all, k_all, v_all, window_size, text_length, has_text=True, dit_seq_shape='30x48x80'): - major, minor = torch.cuda.get_device_capability(q_all.device) - if major == 9 and minor == 0 and sta_fwd is not None: - return sliding_tile_attention_SM90(q_all, k_all, v_all, window_size, text_length, has_text, dit_seq_shape) - elif sliding_tile_attention_triton is not None: - return sliding_tile_attention_triton(q_all, k_all, v_all, window_size, text_length, has_text, dit_seq_shape) - else: - raise ImportError("No suitable sliding tile attention implementation found.") diff --git a/csrc/attn/sliding_tile_attn/st_attn/st_attn_h100.cu b/csrc/attn/sliding_tile_attn/st_attn/st_attn_h100.cu deleted file mode 100644 index 208dc6ecc..000000000 --- a/csrc/attn/sliding_tile_attn/st_attn/st_attn_h100.cu +++ /dev/null @@ -1,841 +0,0 @@ -// # 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)}; - - 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) { - if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 3) { - - cudaFuncSetAttribute( - fwd_attend_ker<128, false, false, true, 1, 1, 1, 5, 6, 10>, - cudaFuncAttributeMaxDynamicSharedMemorySize, - mem_size - ); - fwd_attend_ker<128, false, false, true, 1, 1, 1, 5, 6, 10><<>>(g); - - } else if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 5) { - cudaFuncSetAttribute( - fwd_attend_ker<128, false, false, true, 1, 1, 2, 5, 6, 10>, - cudaFuncAttributeMaxDynamicSharedMemorySize, - mem_size - ); - fwd_attend_ker<128, false, false, true,1, 1, 2, 5, 6, 10><<>>(g); - - } else if (kernel_t_size == 5 && kernel_h_size == 3 && kernel_w_size == 3) { - cudaFuncSetAttribute( - fwd_attend_ker<128, false, false, true, 2, 1, 1, 5, 6, 10>, - cudaFuncAttributeMaxDynamicSharedMemorySize, - mem_size - ); - fwd_attend_ker<128, false, false, true, 2, 1, 1, 5, 6, 10><<>>(g); - - }else if (kernel_t_size ==3 && kernel_h_size == 5 && kernel_w_size == 5){ - cudaFuncSetAttribute( - fwd_attend_ker<128, false, false, true, 1, 2, 2, 5, 6, 10>, - cudaFuncAttributeMaxDynamicSharedMemorySize, - mem_size - ); - fwd_attend_ker<128, false, false, true, 1, 2, 2, 5, 6, 10><<>>(g); - - } else if (kernel_t_size ==5 && kernel_h_size == 6 && kernel_w_size == 1){ - cudaFuncSetAttribute( - fwd_attend_ker<128, false, false, true, 2, 3, 0, 5, 6, 10>, - cudaFuncAttributeMaxDynamicSharedMemorySize, - mem_size - ); - fwd_attend_ker<128, false, false, true, 2, 3, 0, 5, 6, 10><<>>(g); - - } else if (kernel_t_size ==5 && kernel_h_size == 3 && kernel_w_size == 5){ - cudaFuncSetAttribute( - fwd_attend_ker<128, false, false, true, 2, 1, 2, 5, 6, 10>, - cudaFuncAttributeMaxDynamicSharedMemorySize, - mem_size - ); - fwd_attend_ker<128, false, false, true, 2, 1, 2, 5, 6, 10><<>>(g); - - } else if (kernel_t_size == 5 && kernel_h_size == 5 && kernel_w_size == 5){ - cudaFuncSetAttribute( - fwd_attend_ker<128, false, false, true, 2, 2, 2, 5, 6, 10>, - cudaFuncAttributeMaxDynamicSharedMemorySize, - mem_size - ); - fwd_attend_ker<128, false, false, true, 2, 2, 2, 5, 6, 10><<>>(g); - - } else if (kernel_t_size == 5 && kernel_h_size == 5 && kernel_w_size == 7){ - cudaFuncSetAttribute( - fwd_attend_ker<128, false, false, true, 2, 2, 3, 5, 6, 10>, - cudaFuncAttributeMaxDynamicSharedMemorySize, - mem_size - ); - fwd_attend_ker<128, false, false, true, 2, 2, 3, 5, 6, 10><<>>(g); - } else if (kernel_t_size == 5 && kernel_h_size == 6 && kernel_w_size == 10){ - cudaFuncSetAttribute( - fwd_attend_ker<128, false, false, true, 2, 3, 5, 5, 6, 10>, - cudaFuncAttributeMaxDynamicSharedMemorySize, - mem_size - ); - fwd_attend_ker<128, false, false, true, 2, 3, 5, 5, 6, 10><<>>(g); - } else if (kernel_t_size == 5 && kernel_h_size == 1 && kernel_w_size == 1){ - cudaFuncSetAttribute( - fwd_attend_ker<128, false, false, true, 2, 0, 0, 5, 6, 10>, - cudaFuncAttributeMaxDynamicSharedMemorySize, - mem_size - ); - fwd_attend_ker<128, false, false, true, 2, 0, 0, 5, 6, 10><<>>(g); - } else if (kernel_t_size == 1 && kernel_h_size == 6 && kernel_w_size == 10){ - cudaFuncSetAttribute( - fwd_attend_ker<128, false, false, true, 0, 3, 5, 5, 6, 10>, - cudaFuncAttributeMaxDynamicSharedMemorySize, - mem_size - ); - fwd_attend_ker<128, false, false, true, 0, 3, 5, 5, 6, 10><<>>(g); - } else if (kernel_t_size == 5 && kernel_h_size == 1 && kernel_w_size == 10){ - cudaFuncSetAttribute( - fwd_attend_ker<128, false, false, true, 2, 0, 5, 5, 6, 10>, - cudaFuncAttributeMaxDynamicSharedMemorySize, - mem_size - ); - fwd_attend_ker<128, false, false, true,2, 0, 5, 5, 6, 10><<>>(g); - } else { - // print error - std::cout << "Invalid kernel size" << std::endl; - //print kernel size - std::cout << "Kernel size: " << kernel_t_size << " " << kernel_h_size << " " << kernel_w_size << std::endl; - } - } 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){ - if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 3) { - cudaFuncSetAttribute( - fwd_attend_ker<128, false, false, false, 1, 1, 1, 6, 6, 6>, - cudaFuncAttributeMaxDynamicSharedMemorySize, - mem_size - ); - fwd_attend_ker<128, false, false, false, 1, 1, 1, 6, 6, 6><<>>(g); - - } else if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 6) { - cudaFuncSetAttribute( - fwd_attend_ker<128, false, false, false, 1, 1, 3, 6, 6, 6>, - cudaFuncAttributeMaxDynamicSharedMemorySize, - mem_size - ); - fwd_attend_ker<128, false, false, false,1, 1, 3, 6, 6, 6><<>>(g); - - } else if (kernel_t_size == 6 && kernel_h_size == 3 && kernel_w_size == 3) { - cudaFuncSetAttribute( - fwd_attend_ker<128, false, false, false, 3, 1, 1, 6, 6, 6>, - cudaFuncAttributeMaxDynamicSharedMemorySize, - mem_size - ); - fwd_attend_ker<128, false, false, false, 3, 1, 1, 6, 6, 6><<>>(g); - - } else if (kernel_t_size ==3 && kernel_h_size == 6 && kernel_w_size == 6){ - cudaFuncSetAttribute( - fwd_attend_ker<128, false, false, false, 1, 3, 3, 6, 6, 6>, - cudaFuncAttributeMaxDynamicSharedMemorySize, - mem_size - ); - fwd_attend_ker<128, false, false, false, 1, 3, 3, 6, 6, 6><<>>(g); - - }else if (kernel_t_size ==3 && kernel_h_size == 6 && kernel_w_size == 3){ - cudaFuncSetAttribute( - fwd_attend_ker<128, false, false, false, 1, 3, 1, 6, 6, 6>, - cudaFuncAttributeMaxDynamicSharedMemorySize, - mem_size - ); - fwd_attend_ker<128, false, false, false, 1, 3, 1, 6, 6, 6><<>>(g); - - } else if (kernel_t_size ==6 && kernel_h_size == 3 && kernel_w_size == 6){ - cudaFuncSetAttribute( - fwd_attend_ker<128, false, false, false, 3, 1, 3, 6, 6, 6>, - cudaFuncAttributeMaxDynamicSharedMemorySize, - mem_size - ); - fwd_attend_ker<128, false, false, false, 3, 1, 3, 6, 6, 6><<>>(g); - - } else if (kernel_t_size == 6 && kernel_h_size == 6 && kernel_w_size == 6){ - cudaFuncSetAttribute( - fwd_attend_ker<128, false, false, false, 3, 3, 3, 6, 6, 6>, - cudaFuncAttributeMaxDynamicSharedMemorySize, - mem_size - ); - fwd_attend_ker<128, false, false, false, 3, 3, 3, 6, 6, 6><<>>(g); - } else if (kernel_t_size == 6 && kernel_h_size == 1 && kernel_w_size == 1){ - cudaFuncSetAttribute( - fwd_attend_ker<128, false, false, false, 3, 0, 0, 6, 6, 6>, - cudaFuncAttributeMaxDynamicSharedMemorySize, - mem_size - ); - fwd_attend_ker<128, false, false, false, 3, 0, 0, 6, 6, 6><<>>(g); - } else if (kernel_t_size == 6 && kernel_h_size == 1 && kernel_w_size == 6){ - cudaFuncSetAttribute( - fwd_attend_ker<128, false, false, false, 3, 0, 3, 6, 6, 6>, - cudaFuncAttributeMaxDynamicSharedMemorySize, - mem_size - ); - fwd_attend_ker<128, false, false, false, 3, 0, 3, 6, 6, 6><<>>(g); - } else if (kernel_t_size == 6 && kernel_h_size == 6 && kernel_w_size == 1){ - cudaFuncSetAttribute( - fwd_attend_ker<128, false, false, false, 3, 3, 0, 6, 6, 6>, - cudaFuncAttributeMaxDynamicSharedMemorySize, - mem_size - ); - fwd_attend_ker<128, false, false, false, 3, 3, 0, 6, 6, 6><<>>(g); - } else if (kernel_t_size == 1 && kernel_h_size == 6 && kernel_w_size == 6){ - cudaFuncSetAttribute( - fwd_attend_ker<128, false, false, false, 0, 3, 3, 6, 6, 6>, - cudaFuncAttributeMaxDynamicSharedMemorySize, - mem_size - ); - fwd_attend_ker<128, false, false, false, 0, 3, 3, 6, 6, 6><<>>(g); - } else if (kernel_t_size == 1 && kernel_h_size == 1 && kernel_w_size == 6){ - cudaFuncSetAttribute( - fwd_attend_ker<128, false, false, false, 0, 0, 3, 6, 6, 6>, - cudaFuncAttributeMaxDynamicSharedMemorySize, - mem_size - ); - fwd_attend_ker<128, false, false, false, 0, 0, 3, 6, 6, 6><<>>(g); - } else if (kernel_t_size == 1 && kernel_h_size == 6 && kernel_w_size == 1){ - cudaFuncSetAttribute( - fwd_attend_ker<128, false, false, false, 0, 3, 0, 6, 6, 6>, - cudaFuncAttributeMaxDynamicSharedMemorySize, - mem_size - ); - fwd_attend_ker<128, false, false, false, 0, 3, 0, 6, 6, 6><<>>(g); - } else if (kernel_t_size == 6 && kernel_h_size == 6 && kernel_w_size == 1){ - cudaFuncSetAttribute( - fwd_attend_ker<128, false, false, false, 3, 3, 0, 6, 6, 6>, - cudaFuncAttributeMaxDynamicSharedMemorySize, - mem_size - ); - fwd_attend_ker<128, false, false, false, 3, 3, 0, 6, 6, 6><<>>(g); - } else if (kernel_t_size == 6 && kernel_h_size == 1 && kernel_w_size == 6){ - cudaFuncSetAttribute( - fwd_attend_ker<128, false, false, false, 3, 0, 3, 6, 6, 6>, - cudaFuncAttributeMaxDynamicSharedMemorySize, - mem_size - ); - fwd_attend_ker<128, false, false, false, 3, 0, 3, 6, 6, 6><<>>(g); - } else { - // print error - std::cout << "Invalid kernel size" << std::endl; - //print kernel size - std::cout << "Kernel size: " << kernel_t_size << " " << kernel_h_size << " " << kernel_w_size << std::endl; - } - } - else if (kernel_aspect_ratio_flag == 3) { - if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 3) { - - cudaFuncSetAttribute( - fwd_attend_ker<128, false, false, false, 1, 1, 1, 3, 6, 10>, - cudaFuncAttributeMaxDynamicSharedMemorySize, - mem_size - ); - fwd_attend_ker<128, false, false, false, 1, 1, 1, 3, 6, 10><<>>(g); - - } else if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 5) { - cudaFuncSetAttribute( - fwd_attend_ker<128, false, false, false, 1, 1, 2, 3, 6, 10>, - cudaFuncAttributeMaxDynamicSharedMemorySize, - mem_size - ); - fwd_attend_ker<128, false, false, false,1, 1, 2, 3, 6, 10><<>>(g); - - } else if (kernel_t_size == 3 && kernel_h_size == 5 && kernel_w_size == 5){ - cudaFuncSetAttribute( - fwd_attend_ker<128, false, false, false, 1, 2, 2, 3, 6, 10>, - cudaFuncAttributeMaxDynamicSharedMemorySize, - mem_size - ); - fwd_attend_ker<128, false, false, false, 1, 2, 2, 3, 6, 10><<>>(g); - - } else if (kernel_t_size == 3 && kernel_h_size == 6 && kernel_w_size == 1){ - cudaFuncSetAttribute( - fwd_attend_ker<128, false, false, false, 1, 3, 0, 3, 6, 10>, - cudaFuncAttributeMaxDynamicSharedMemorySize, - mem_size - ); - fwd_attend_ker<128, false, false, false, 1, 3, 0, 3, 6, 10><<>>(g); - - } else if (kernel_t_size == 3 && kernel_h_size == 5 && kernel_w_size == 7){ - cudaFuncSetAttribute( - fwd_attend_ker<128, false, false, false, 1, 2, 3, 3, 6, 10>, - cudaFuncAttributeMaxDynamicSharedMemorySize, - mem_size - ); - fwd_attend_ker<128, false, false, false, 1, 2, 3, 3, 6, 10><<>>(g); - } else if (kernel_t_size == 3 && kernel_h_size == 5 && kernel_w_size == 9){ - cudaFuncSetAttribute( - fwd_attend_ker<128, false, false, false, 1, 2, 4, 3, 6, 10>, - cudaFuncAttributeMaxDynamicSharedMemorySize, - mem_size - ); - fwd_attend_ker<128, false, false, false, 1, 2, 4, 3, 6, 10><<>>(g); - } else if (kernel_t_size == 3 && kernel_h_size == 6 && kernel_w_size == 10){ - cudaFuncSetAttribute( - fwd_attend_ker<128, false, false, false, 1, 3, 5, 3, 6, 10>, - cudaFuncAttributeMaxDynamicSharedMemorySize, - mem_size - ); - fwd_attend_ker<128, false, false, false, 1, 3, 5, 3, 6, 10><<>>(g); - } else if (kernel_t_size == 3 && kernel_h_size == 6 && kernel_w_size == 3){ - cudaFuncSetAttribute( - fwd_attend_ker<128, false, false, false, 1, 3, 1, 3, 6, 10>, - cudaFuncAttributeMaxDynamicSharedMemorySize, - mem_size - ); - fwd_attend_ker<128, false, false, false, 1, 3, 1, 3, 6, 10><<>>(g); - } else if (kernel_t_size == 3 && kernel_h_size == 1 && kernel_w_size == 1){ - cudaFuncSetAttribute( - fwd_attend_ker<128, false, false, false, 1, 0, 0, 3, 6, 10>, - cudaFuncAttributeMaxDynamicSharedMemorySize, - mem_size - ); - fwd_attend_ker<128, false, false, false, 1, 0, 0, 3, 6, 10><<>>(g); - } else if (kernel_t_size == 1 && kernel_h_size == 6 && kernel_w_size == 10){ - cudaFuncSetAttribute( - fwd_attend_ker<128, false, false, false, 0, 3, 5, 3, 6, 10>, - cudaFuncAttributeMaxDynamicSharedMemorySize, - mem_size - ); - fwd_attend_ker<128, false, false, false, 0, 3, 5, 3, 6, 10><<>>(g); - } else if (kernel_t_size == 1 && kernel_h_size == 5 && kernel_w_size == 10){ - cudaFuncSetAttribute( - fwd_attend_ker<128, false, false, false, 0, 2, 5, 3, 6, 10>, - cudaFuncAttributeMaxDynamicSharedMemorySize, - mem_size - ); - fwd_attend_ker<128, false, false, false, 0, 2, 5, 3, 6, 10><<>>(g); - } else if (kernel_t_size == 1 && kernel_h_size == 6 && kernel_w_size == 7){ - cudaFuncSetAttribute( - fwd_attend_ker<128, false, false, false, 0, 3, 3, 3, 6, 10>, - cudaFuncAttributeMaxDynamicSharedMemorySize, - mem_size - ); - fwd_attend_ker<128, false, false, false, 0, 3, 3, 3, 6, 10><<>>(g); - } else if (kernel_t_size == 1 && kernel_h_size == 5 && kernel_w_size == 7){ - cudaFuncSetAttribute( - fwd_attend_ker<128, false, false, false, 0, 2, 3, 3, 6, 10>, - cudaFuncAttributeMaxDynamicSharedMemorySize, - mem_size - ); - fwd_attend_ker<128, false, false, false, 0, 2, 3, 3, 6, 10><<>>(g); - } else if (kernel_t_size == 1 && kernel_h_size == 5 && kernel_w_size == 9){ - cudaFuncSetAttribute( - fwd_attend_ker<128, false, false, false, 0, 2, 4, 3, 6, 10>, - cudaFuncAttributeMaxDynamicSharedMemorySize, - mem_size - ); - fwd_attend_ker<128, false, false, false, 0, 2, 4, 3, 6, 10><<>>(g); - } else if (kernel_t_size == 3 && kernel_h_size == 1 && kernel_w_size == 10){ - cudaFuncSetAttribute( - fwd_attend_ker<128, false, false, false, 1, 0, 5, 3, 6, 10>, - cudaFuncAttributeMaxDynamicSharedMemorySize, - mem_size - ); - fwd_attend_ker<128, false, false, false,1, 0, 5, 3, 6, 10><<>>(g); - } else if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 10){ - cudaFuncSetAttribute( - fwd_attend_ker<128, false, false, false, 1, 1, 5, 3, 6, 10>, - cudaFuncAttributeMaxDynamicSharedMemorySize, - mem_size - ); - fwd_attend_ker<128, false, false, false,1, 1, 5, 3, 6, 10><<>>(g); - } else if (kernel_t_size == 1 && kernel_h_size == 3 && kernel_w_size == 10){ - cudaFuncSetAttribute( - fwd_attend_ker<128, false, false, false, 0, 1, 5, 3, 6, 10>, - cudaFuncAttributeMaxDynamicSharedMemorySize, - mem_size - ); - fwd_attend_ker<128, false, false, false,0, 1, 5, 3, 6, 10><<>>(g); - } else if (kernel_t_size == 1 && kernel_h_size == 6 && kernel_w_size == 5){ - cudaFuncSetAttribute( - fwd_attend_ker<128, false, false, false, 0, 3, 2, 3, 6, 10>, - cudaFuncAttributeMaxDynamicSharedMemorySize, - mem_size - ); - fwd_attend_ker<128, false, false, false,0, 3, 2, 3, 6, 10><<>>(g); - } else { - // print error - std::cout << "Invalid kernel size" << std::endl; - //print kernel size - std::cout << "Kernel size: " << kernel_t_size << " " << kernel_h_size << " " << kernel_w_size << std::endl; - } - } - - else { - std::cout << "Unsupported kernel_aspect_ratio_flag: " << kernel_aspect_ratio_flag << std::endl; - } - - } - CHECK_CUDA_ERROR(cudaGetLastError()); - // cudaStreamSynchronize(stream); - } - - return o; - //cudadevicesynchronize(); -} - diff --git a/csrc/attn/tests/flex_sta_ref.py b/csrc/attn/tests/flex_sta_ref.py deleted file mode 100644 index d7d9528a6..000000000 --- a/csrc/attn/tests/flex_sta_ref.py +++ /dev/null @@ -1,71 +0,0 @@ -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/attn/tests/test_sta.py b/csrc/attn/tests/test_sta.py deleted file mode 100644 index 699e56cd6..000000000 --- a/csrc/attn/tests/test_sta.py +++ /dev/null @@ -1,87 +0,0 @@ -import torch -from flex_sta_ref import get_sliding_tile_attention_mask -from st_attn import sliding_tile_attention -from torch.nn.attention.flex_attention import flex_attention -# from flash_attn_interface import flash_attn_func -from tqdm import tqdm - -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): - 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=50, error_mode='all'): - results = { - 'TK vs FLEX': { - 'sum_diff': 0, - 'sum_abs': 0, - 'max_diff': 0 - }, - } - kernel_size_ls = [(3, 3, 5), (3, 1, 10)] - from tqdm import tqdm - for kernel_size in tqdm(kernel_size_ls): - for _ in range(num_iterations): - torch.manual_seed(0) - - 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) - results['TK vs FLEX']['sum_diff'] += torch.sum(abs_diff).item() - results['TK vs FLEX']['max_diff'] = max(results['TK vs FLEX']['max_diff'], torch.max(abs_diff).item()) - - torch.cuda.empty_cache() - print("kernel_size", kernel_size) - print("max_diff", torch.max(abs_diff).item()) - print( - "avg_diff", - torch.sum(abs_diff).item() / (b * h * n * d * - (1 if error_mode == 'output' else 3 if error_mode == 'backward' else 4))) - - total_elements = b * h * n * d * num_iterations * (1 if error_mode == 'output' else - 3 if error_mode == 'backward' else 4) * len(kernel_size_ls) - for name, data in results.items(): - avg_diff = data['sum_diff'] / total_elements - max_diff = data['max_diff'] - results[name] = {'avg_diff': avg_diff, 'max_diff': max_diff} - - return results - - -# Example usage -b, h, d = 2, 24, 128 -n = 69120 # Sequence length -causal = False -mean = 1e-1 -std = 10 - -# Run correctness check directly -results = check_correctness(b, h, n, d, causal, mean, std, error_mode='output') -assert results['TK vs FLEX']['avg_diff'] < 3e-6, f"Average difference: {results['TK vs FLEX']['avg_diff']} is too large" -assert results['TK vs FLEX']['max_diff'] < 4e-2, f"Maximum difference: {results['TK vs FLEX']['max_diff']} is too large" -print(f"Average difference: {results['TK vs FLEX']['avg_diff']}") -print(f"Maximum difference: {results['TK vs FLEX']['max_diff']}") diff --git a/csrc/attn/tests/test_vsa.py b/csrc/attn/tests/test_vsa.py deleted file mode 100644 index e2da3ec58..000000000 --- a/csrc/attn/tests/test_vsa.py +++ /dev/null @@ -1,245 +0,0 @@ -import torch -import sys -import os -import numpy as np -from tqdm import tqdm - -# Add the parent directory to the path to import block_sparse_attn -sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) -from tests.utils import generate_block_sparse_mask_for_function, create_full_mask_from_block_mask -from vsa import block_sparse_attn - -BLOCK_M = 64 -BLOCK_N = 64 - -def pytorch_test(Q, K, V, block_sparse_mask, dO): - q_ = Q.clone().float().requires_grad_() - k_ = K.clone().float().requires_grad_() - v_ = V.clone().float().requires_grad_() - - QK = torch.matmul(q_, k_.transpose(-2, -1)) - QK /= (q_.size(-1) ** 0.5) - QK = QK.masked_fill(~block_sparse_mask.unsqueeze(0), float('-inf')) - - QK = torch.nn.functional.softmax(QK, dim=-1) - output = torch.matmul(QK, v_) - - dO_ = dO - output.backward(dO_) - return ( - output.to(torch.bfloat16), - q_.grad.to(torch.bfloat16), - k_.grad.to(torch.bfloat16), - v_.grad.to(torch.bfloat16), - ) - - -def block_sparse_kernel_test(Q, K, V, block_sparse_mask, variable_block_sizes, q_non_pad_index, kv_non_pad_index, q_num_blocks, kv_num_blocks, dO): - Q = Q.detach().requires_grad_() - K = K.detach().requires_grad_() - V = V.detach().requires_grad_() - - q_padded = vsa_pad(Q, q_non_pad_index, q_num_blocks, BLOCK_M) - k_padded = vsa_pad(K, kv_non_pad_index, kv_num_blocks, BLOCK_M) - v_padded = vsa_pad(V, kv_non_pad_index, kv_num_blocks, BLOCK_M) - output, _= block_sparse_attn(q_padded, k_padded, v_padded, block_sparse_mask, variable_block_sizes) - output = output[:, :, q_non_pad_index, :] - output.backward(dO) - return output, Q.grad, K.grad, V.grad - - -def get_non_pad_index( - vid_len: torch.LongTensor, - n_win: int, - win_size: int, -): - device = vid_len.device - starts_pad = torch.arange(n_win, device=device) * win_size - index_pad = starts_pad[:, None] + torch.arange(win_size, device=device)[None, :] - index_mask = torch.arange(win_size, device=device)[None, :] < vid_len[:, None] - - return index_pad[index_mask] - -def generate_tensor(shape, dtype, device): - tensor = torch.randn(shape, dtype=dtype, device=device) - return tensor - -def generate_variable_block_sizes(num_blocks, min_size=16, max_size=64, device="cuda"): - return torch.randint(min_size, max_size + 1, (num_blocks,), device=device, dtype=torch.int32) - - -def vsa_pad(x, non_pad_index, num_blocks, block_size): - padded_x = torch.zeros((1, x.shape[1], num_blocks * BLOCK_M, x.shape[3]), device=x.device, dtype=x.dtype) - padded_x[:, :, non_pad_index, :] = x - return padded_x - -def check_correctness(h, d, num_blocks, k, num_iterations=20, error_mode='all'): - results = { - 'gO': {'sum_diff': 0.0, 'sum_abs': 0.0, 'max_diff': 0.0}, - 'gQ': {'sum_diff': 0.0, 'sum_abs': 0.0, 'max_diff': 0.0}, - 'gK': {'sum_diff': 0.0, 'sum_abs': 0.0, 'max_diff': 0.0}, - 'gV': {'sum_diff': 0.0, 'sum_abs': 0.0, 'max_diff': 0.0}, - } - - device = "cuda" if torch.cuda.is_available() else "cpu" - variable_block_sizes = generate_variable_block_sizes(num_blocks, device=device) - S = int(variable_block_sizes.sum().item()) - padded_S = num_blocks * BLOCK_M - non_pad_index = get_non_pad_index(variable_block_sizes, num_blocks, BLOCK_M) - block_mask = generate_block_sparse_mask_for_function(h, num_blocks, num_blocks, k, device) - full_mask = create_full_mask_from_block_mask(block_mask, variable_block_sizes, variable_block_sizes, device) - for _ in range(num_iterations): - Q = generate_tensor((1, h, S, d), torch.bfloat16, device) - K = generate_tensor((1, h, S, d), torch.bfloat16, device) - V = generate_tensor((1, h, S, d), torch.bfloat16, device) - dO = generate_tensor((1, h, S, d), torch.bfloat16, device) - - # print(Q.shape, K.shape, V.shape, dO.shape) - - # dO_padded = torch.zeros_like(dO_padded) - # dO_padded[:, :, non_pad_index, :] = dO - - pt_o, pt_qg, pt_kg, pt_vg = pytorch_test(Q, K, V, full_mask, dO) - bs_o, bs_qg, bs_kg, bs_vg = block_sparse_kernel_test(Q, K, V, block_mask.unsqueeze(0), variable_block_sizes, non_pad_index, non_pad_index, num_blocks, num_blocks, dO) - for name, (pt, bs) in zip(['gQ', 'gK', 'gV', 'gO'], [(pt_qg, bs_qg), (pt_kg, bs_kg), (pt_vg, bs_vg), (pt_o, bs_o)]): - if bs is not None: - diff = pt - bs - abs_diff = torch.abs(diff) - results[name]['sum_diff'] += torch.sum(abs_diff).item() - results[name]['sum_abs'] += torch.sum(torch.abs(pt)).item() - rel_max_diff = torch.max(abs_diff) / torch.mean(torch.abs(pt)) - results[name]['max_diff'] = max(results[name]['max_diff'], rel_max_diff.item()) - if torch.cuda.is_available(): - torch.cuda.empty_cache() - - total_elements = h * S * d * num_iterations - for name, data in results.items(): - avg_diff = data['sum_diff'] / total_elements - max_diff = data['max_diff'] - results[name] = {'avg_diff': avg_diff, 'max_diff': max_diff} - - return results - -def check_correctness_qkdiff(h, d, num_q_blocks, num_kv_blocks, k, num_iterations=20, error_mode='all'): - results = { - 'gO': {'sum_diff': 0.0, 'sum_abs': 0.0, 'max_diff': 0.0}, - 'gQ': {'sum_diff': 0.0, 'sum_abs': 0.0, 'max_diff': 0.0}, - 'gK': {'sum_diff': 0.0, 'sum_abs': 0.0, 'max_diff': 0.0}, - 'gV': {'sum_diff': 0.0, 'sum_abs': 0.0, 'max_diff': 0.0}, - } - - device = "cuda" if torch.cuda.is_available() else "cpu" - q_variable_block_sizes = generate_variable_block_sizes(num_q_blocks, device=device) - kv_variable_block_sizes = generate_variable_block_sizes(num_kv_blocks, device=device) - - S_q = int(q_variable_block_sizes.sum().item()) - S_kv = int(kv_variable_block_sizes.sum().item()) - - q_non_pad_index = get_non_pad_index(q_variable_block_sizes, num_q_blocks, BLOCK_M) - kv_non_pad_index = get_non_pad_index(kv_variable_block_sizes, num_kv_blocks, BLOCK_M) - - block_mask = generate_block_sparse_mask_for_function(h, num_q_blocks, num_kv_blocks, k, device) - full_mask = create_full_mask_from_block_mask(block_mask, q_variable_block_sizes, kv_variable_block_sizes, device) - - for _ in range(num_iterations): - Q = generate_tensor((1, h, S_q, d), torch.bfloat16, device) - K = generate_tensor((1, h, S_kv, d), torch.bfloat16, device) - V = generate_tensor((1, h, S_kv, d), torch.bfloat16, device) - dO = generate_tensor((1, h, S_q, d), torch.bfloat16, device) - - # print(Q.shape, K.shape, V.shape, dO.shape) - - pt_o, pt_qg, pt_kg, pt_vg = pytorch_test(Q, K, V, full_mask, dO) - bs_o, bs_qg, bs_kg, bs_vg = block_sparse_kernel_test(Q, K, V, block_mask.unsqueeze(0), kv_variable_block_sizes, q_non_pad_index, kv_non_pad_index, num_q_blocks, num_kv_blocks, dO) - - for name, (pt, bs) in zip(['gQ', 'gK', 'gV', 'gO'], [(pt_qg, bs_qg), (pt_kg, bs_kg), (pt_vg, bs_vg), (pt_o, bs_o)]): - if bs is not None: - diff = pt - bs - abs_diff = torch.abs(diff) - results[name]['sum_diff'] += torch.sum(abs_diff).item() - results[name]['sum_abs'] += torch.sum(torch.abs(pt)).item() - rel_max_diff = torch.max(abs_diff) / torch.mean(torch.abs(pt)) - results[name]['max_diff'] = max(results[name]['max_diff'], rel_max_diff.item()) - if torch.cuda.is_available(): - torch.cuda.empty_cache() - - total_elements_q = h * S_q * d * num_iterations - total_elements_kv = h * S_kv * d * num_iterations - - for name, data in results.items(): - total_elements = total_elements_q if name in ['gQ', 'gO'] else total_elements_kv - avg_diff = data['sum_diff'] / total_elements - max_diff = data['max_diff'] - results[name] = {'avg_diff': avg_diff, 'max_diff': max_diff} - - return results - -def generate_error_graphs(h, d, error_mode='all'): - test_configs = [ - {"num_blocks": 16, "k": 2, "description": "Small sequence"}, - {"num_blocks": 32, "k": 4, "description": "Medium sequence"}, - {"num_blocks": 53, "k": 6, "description": "Large sequence"}, - ] - - print(f"\nError Analysis for h={h}, d={d}, mode={error_mode}") - print("=" * 150) - print(f"{'Config':<20} {'Blocks':<8} {'K':<4} " - f"{'gQ Avg':<12} {'Rel gQ Max':<12} " - f"{'gK Avg':<12} {'Rel gK Max':<12} " - f"{'gV Avg':<12} {'Rel gV Max':<12} " - f"{'gO Avg':<12} {'Rel gO Max':<12}") - print("-" * 150) - - for config in test_configs: - num_blocks = config["num_blocks"] - k = config["k"] - description = config["description"] - results = check_correctness(h, d, num_blocks, k, error_mode=error_mode) - print(f"{description:<20} {num_blocks:<8} {k:<4} " - f"{results['gQ']['avg_diff']:<12.6e} {results['gQ']['max_diff']:<12.6e} " - f"{results['gK']['avg_diff']:<12.6e} {results['gK']['max_diff']:<12.6e} " - f"{results['gV']['avg_diff']:<12.6e} {results['gV']['max_diff']:<12.6e} " - f"{results['gO']['avg_diff']:<12.6e} {results['gO']['max_diff']:<12.6e}") - - print("-" * 150) - -def generate_error_graphs_qkdiff(h, d, error_mode='all'): - test_configs = [ - {"num_q_blocks": 16, "num_kv_blocks": 32, "k": 2, "description": "Small Q, Med KV"}, - {"num_q_blocks": 32, "num_kv_blocks": 16, "k": 4, "description": "Med Q, Small KV"}, - {"num_q_blocks": 53, "num_kv_blocks": 32, "k": 6, "description": "Large Q, Med KV"}, - {"num_q_blocks": 16, "num_kv_blocks": 48, "k": 2, "description": "Small Q, Large KV"}, - {"num_q_blocks": 48, "num_kv_blocks": 16, "k": 2, "description": "Large Q, Small KV"}, - ] - - print(f"\nError Analysis (QK Diff) for h={h}, d={d}, mode={error_mode}") - print("=" * 150) - print(f"{'Config':<20} {'Q Blks':<8} {'KV Blks':<8} {'K':<4} " - f"{'gQ Avg':<12} {'Rel gQ Max':<12} " - f"{'gK Avg':<12} {'Rel gK Max':<12} " - f"{'gV Avg':<12} {'Rel gV Max':<12} " - f"{'gO Avg':<12} {'Rel gO Max':<12}") - print("-" * 150) - - for config in test_configs: - num_q_blocks = config["num_q_blocks"] - num_kv_blocks = config["num_kv_blocks"] - k = config["k"] - description = config["description"] - results = check_correctness_qkdiff(h, d, num_q_blocks, num_kv_blocks, k, error_mode=error_mode) - print(f"{description:<20} {num_q_blocks:<8} {num_kv_blocks:<8} {k:<4} " - f"{results['gQ']['avg_diff']:<12.6e} {results['gQ']['max_diff']:<12.6e} " - f"{results['gK']['avg_diff']:<12.6e} {results['gK']['max_diff']:<12.6e} " - f"{results['gV']['avg_diff']:<12.6e} {results['gV']['max_diff']:<12.6e} " - f"{results['gO']['avg_diff']:<12.6e} {results['gO']['max_diff']:<12.6e}") - - print("-" * 150) - -if __name__ == "__main__": - h, d = 16, 128 - print("Block Sparse Attention with Variable Block Sizes Analysis") - print("=" * 60) - for mode in ['backward']: - generate_error_graphs(h, d, error_mode=mode) - generate_error_graphs_qkdiff(h, d, error_mode=mode) - print("\nAnalysis completed for all modes.") \ No newline at end of file diff --git a/csrc/attn/tests/test_vsa_forward.py b/csrc/attn/tests/test_vsa_forward.py deleted file mode 100644 index d20055094..000000000 --- a/csrc/attn/tests/test_vsa_forward.py +++ /dev/null @@ -1,236 +0,0 @@ -import os -import sys -from typing import Tuple - -import torch - -# Make sure we can import from the project root (`vsa`, `tests.utils`, etc.) -CURRENT_DIR = os.path.dirname(os.path.abspath(__file__)) -PROJECT_ROOT = os.path.dirname(CURRENT_DIR) -if PROJECT_ROOT not in sys.path: - sys.path.append(PROJECT_ROOT) -if CURRENT_DIR not in sys.path: - sys.path.append(CURRENT_DIR) - -from tests.utils import ( - generate_block_sparse_mask_for_function, - create_full_mask_from_block_mask, -) -from vsa import block_sparse_attn, BLOCK_M -import test_vsa as ref # reuse helper functions from backward test - - -def pytorch_forward( - Q: torch.Tensor, - K: torch.Tensor, - V: torch.Tensor, - block_sparse_mask: torch.Tensor, -) -> torch.Tensor: - """ - Dense PyTorch reference forward: - - Q: [1, h, S_q, d] - - K,V: [1, h, S_kv, d] - - block_sparse_mask: [h, S_q, S_kv] bool - """ - q = Q.clone().float() - k = K.clone().float() - v = V.clone().float() - - attn = torch.matmul(q, k.transpose(-2, -1)) # [1, h, S_q, S_kv] - attn = attn / (q.size(-1) ** 0.5) - attn = attn.masked_fill(~block_sparse_mask.unsqueeze(0), float("-inf")) - attn = torch.nn.functional.softmax(attn, dim=-1) - out = torch.matmul(attn, v) # [1, h, S_q, d] - return out.to(torch.bfloat16) - - -def block_sparse_forward_test( - Q: torch.Tensor, - K: torch.Tensor, - V: torch.Tensor, - block_sparse_mask: torch.Tensor, - variable_block_sizes: torch.Tensor, - q_non_pad_index: torch.Tensor, - kv_non_pad_index: torch.Tensor, - q_num_blocks: int, - kv_num_blocks: int, -) -> torch.Tensor: - """ - Forward-only wrapper around `block_sparse_attn`, mirroring `block_sparse_kernel_test` - but without any backward / grad logic. - """ - Q = Q.detach() - K = K.detach() - V = V.detach() - - q_padded = ref.vsa_pad(Q, q_non_pad_index, q_num_blocks, BLOCK_M) - k_padded = ref.vsa_pad(K, kv_non_pad_index, kv_num_blocks, BLOCK_M) - v_padded = ref.vsa_pad(V, kv_non_pad_index, kv_num_blocks, BLOCK_M) - - out_padded, _ = block_sparse_attn( - q_padded, k_padded, v_padded, block_sparse_mask, variable_block_sizes - ) - # Remove padding on the query side - out = out_padded[:, :, q_non_pad_index, :] - return out - - -def run_forward_equal_qk( - h: int = 16, - d: int = 128, - num_blocks: int = 16, - k: int = 2, - num_iterations: int = 5, -) -> Tuple[float, float]: - """ - Forward-only correctness test for the case S_q == S_kv. - Mirrors `check_correctness` but only compares forward outputs. - """ - assert torch.cuda.is_available(), "VSA kernels require CUDA" - device = "cuda" - - variable_block_sizes = ref.generate_variable_block_sizes( - num_blocks, device=device - ) - S = int(variable_block_sizes.sum().item()) - non_pad_index = ref.get_non_pad_index( - variable_block_sizes, num_blocks, BLOCK_M - ) - - block_mask = generate_block_sparse_mask_for_function( - h, num_blocks, num_blocks, k, device - ) - full_mask = create_full_mask_from_block_mask( - block_mask, variable_block_sizes, variable_block_sizes, device - ) - print(f"[qkequal] h: {h}, d: {d}, num_blocks: {num_blocks}, k: {k}") - print(f"[qkequal] variable_block_sizes: {variable_block_sizes}, non_pad_index: {non_pad_index.shape}, block_mask: {block_mask.shape}, full_mask: {full_mask.shape}") - sum_diff = 0.0 - sum_abs = 0.0 - max_rel_diff = 0.0 - - for i in range(num_iterations): - Q = ref.generate_tensor((1, h, S, d), torch.bfloat16, device) - K = ref.generate_tensor((1, h, S, d), torch.bfloat16, device) - V = ref.generate_tensor((1, h, S, d), torch.bfloat16, device) - - if i == 0: print(f"[qkequal] Q: {Q.shape}, K: {K.shape}, V: {V.shape}, full_mask: {full_mask.shape}") - if i == 0: print(f"[qkequal] block_mask: {block_mask.shape}") - - pt_o = pytorch_forward(Q, K, V, full_mask) - bs_o = block_sparse_forward_test( - Q, - K, - V, - block_mask.unsqueeze(0), - variable_block_sizes, - non_pad_index, - non_pad_index, - num_blocks, - num_blocks, - ) - - diff = (pt_o - bs_o).abs() - sum_diff += diff.sum().item() - sum_abs += pt_o.abs().sum().item() - rel_max = diff.max() / (pt_o.abs().mean() + 1e-6) - max_rel_diff = max(max_rel_diff, rel_max.item()) - - total_elems = h * S * d * num_iterations - avg_abs_err = sum_diff / total_elems - return avg_abs_err, max_rel_diff - - -def run_forward_qk_diff( - h: int = 16, - d: int = 128, - num_q_blocks: int = 16, - num_kv_blocks: int = 32, - k: int = 2, - num_iterations: int = 5, -) -> Tuple[float, float]: - """ - Forward-only correctness test for the case S_q != S_kv. - - NOTE: - - The Triton backend supports different Q/KV logical lengths via padding. - - The SM90 (H100) CUDA backend currently assumes the same number of blocks - for Q and KV, so we skip this test there. - """ - assert torch.cuda.is_available(), "VSA kernels require CUDA" - - device = "cuda" - - q_variable_block_sizes = ref.generate_variable_block_sizes( - num_q_blocks, device=device - ) - kv_variable_block_sizes = ref.generate_variable_block_sizes( - num_kv_blocks, device=device - ) - - S_q = int(q_variable_block_sizes.sum().item()) - S_kv = int(kv_variable_block_sizes.sum().item()) - - q_non_pad_index = ref.get_non_pad_index( - q_variable_block_sizes, num_q_blocks, BLOCK_M - ) - kv_non_pad_index = ref.get_non_pad_index( - kv_variable_block_sizes, num_kv_blocks, BLOCK_M - ) - - block_mask = generate_block_sparse_mask_for_function( - h, num_q_blocks, num_kv_blocks, k, device - ) - full_mask = create_full_mask_from_block_mask( - block_mask, q_variable_block_sizes, kv_variable_block_sizes, device - ) - - sum_diff = 0.0 - sum_abs = 0.0 - max_rel_diff = 0.0 - - for _ in range(num_iterations): - Q = ref.generate_tensor((1, h, S_q, d), torch.bfloat16, device) - K = ref.generate_tensor((1, h, S_kv, d), torch.bfloat16, device) - V = ref.generate_tensor((1, h, S_kv, d), torch.bfloat16, device) - - pt_o = pytorch_forward(Q, K, V, full_mask) - bs_o = block_sparse_forward_test( - Q, - K, - V, - block_mask.unsqueeze(0), - kv_variable_block_sizes, - q_non_pad_index, - kv_non_pad_index, - num_q_blocks, - num_kv_blocks, - ) - - diff = (pt_o - bs_o).abs() - sum_diff += diff.sum().item() - sum_abs += pt_o.abs().sum().item() - rel_max = diff.max() / (pt_o.abs().mean() + 1e-6) - max_rel_diff = max(max_rel_diff, rel_max.item()) - - total_elems = h * S_q * d * num_iterations - avg_abs_err = sum_diff / total_elems - return avg_abs_err, max_rel_diff - - -if __name__ == "__main__": - h, d = 16, 128 - print("Forward Block Sparse Attention Check (QK Equal)") - print("=" * 80) - avg_err_eq, max_rel_eq = run_forward_equal_qk(h, d, num_blocks=32, k=2) - print(f"QK equal: avg |ΔO| = {avg_err_eq:.6e}, max rel ΔO = {max_rel_eq:.6e}") - - print("\nForward Block Sparse Attention Check (QK Different)") - print("=" * 80) - avg_err_diff, max_rel_diff = run_forward_qk_diff( - h, d, num_q_blocks=32, num_kv_blocks=48, k=2 - ) - print( - f"QK diff: avg |ΔO| = {avg_err_diff:.6e}, max rel ΔO = {max_rel_diff:.6e}" - ) - diff --git a/csrc/attn/tests/utils.py b/csrc/attn/tests/utils.py deleted file mode 100644 index 3d9f15f45..000000000 --- a/csrc/attn/tests/utils.py +++ /dev/null @@ -1,60 +0,0 @@ -import torch - -def generate_block_sparse_mask_for_function(h, num_q_blocks, num_kv_blocks, k, device="cuda"): - """ - Generate block sparse mask of shape [h, num_q_blocks, num_kv_blocks]. - - Args: - h: number of heads - num_q_blocks: number of query blocks - num_kv_blocks: number of key/value blocks - k: number of kv blocks each q block attends to - device: device to create tensors on - - Returns: - block_sparse_mask: [h, num_q_blocks, num_kv_blocks] bool tensor - """ - k = min(k, num_kv_blocks) - scores = torch.rand(h, num_q_blocks, num_kv_blocks, device=device) - _, indices = torch.topk(scores, k, dim=-1) - block_sparse_mask = torch.zeros(h, num_q_blocks, num_kv_blocks, dtype=torch.bool, device=device) - - block_sparse_mask = block_sparse_mask.scatter_(2, indices, 1).bool() - return block_sparse_mask - - -def create_full_mask_from_block_mask(block_sparse_mask, q_variable_block_sizes, - kv_variable_block_sizes, device="cuda"): - """ - Convert block-level sparse mask to full attention mask. - - Args: - block_sparse_mask: [h, num_q_blocks, num_kv_blocks] bool tensor - q_variable_block_sizes: [num_q_blocks] tensor - kv_variable_block_sizes: [num_kv_blocks] tensor - device: device to create tensors on - - Returns: - full_mask: [h, S_q, S_kv] bool tensor where S = total sequence length - """ - h, num_q_blocks, num_kv_blocks = block_sparse_mask.shape - total_q_seq_len = q_variable_block_sizes.sum().item() - total_kv_seq_len = kv_variable_block_sizes.sum().item() - - q_cumsum = torch.cat([torch.tensor([0], device=device), q_variable_block_sizes.cumsum(dim=0)[:-1]]) - kv_cumsum = torch.cat([torch.tensor([0], device=device), kv_variable_block_sizes.cumsum(dim=0)[:-1]]) - - full_mask = torch.zeros(h, total_q_seq_len, total_kv_seq_len, dtype=torch.bool, device=device) - - for head in range(h): - for q_block in range(num_q_blocks): - q_start = q_cumsum[q_block] - q_end = q_start + q_variable_block_sizes[q_block] - - for kv_block in range(num_kv_blocks): - if block_sparse_mask[head, q_block, kv_block]: - kv_start = kv_cumsum[kv_block] - kv_end = kv_start + kv_variable_block_sizes[kv_block] - full_mask[head, q_start:q_end, kv_start:kv_end] = True - - return full_mask \ No newline at end of file diff --git a/csrc/attn/video_sparse_attn/MANIFEST.in b/csrc/attn/video_sparse_attn/MANIFEST.in deleted file mode 100644 index de701fcb9..000000000 --- a/csrc/attn/video_sparse_attn/MANIFEST.in +++ /dev/null @@ -1,2 +0,0 @@ -recursive-include tk * -include config_vsa.py diff --git a/csrc/attn/video_sparse_attn/README.md b/csrc/attn/video_sparse_attn/README.md deleted file mode 100644 index c4a027592..000000000 --- a/csrc/attn/video_sparse_attn/README.md +++ /dev/null @@ -1,61 +0,0 @@ - - -# Attention Kernel Used in FastVideo - -## Video Sparse Attention (VSA) - -### Installation -We support H100 (via TK) and any other GPU (via triton) for VSA. - -```bash -pip install vsa -``` - -Install from source: - -```bash -git submodule update --init --recursive -python setup.py install -``` - - -If you encounter error during installation, try below: -Install C++20 for ThunderKittens: -```bash -sudo apt update -sudo apt install gcc-11 g++-11 - -sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11 - -sudo apt update -sudo apt install clang-11 -``` -(If you use CUDA12.8) -```bash -export CUDA_HOME=/usr/local/cuda-12.8 -export PATH=${CUDA_HOME}/bin:${PATH} -export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH -``` - -### Verify if you have successfully installed - -```bash -# test numerical -python ../tests/test_vsa.py -# (For H100) test speed -python ../benchmarks/bench_vsa_hopper.py -``` - -bench_vsa_hopper.py should print something like this: - -```bash -Using topk=76 kv blocks per q block (out of 768 total kv blocks) - -=== BLOCK SPARSE ATTENTION BENCHMARK === -Block Sparse Forward - TFLOPS: 5622.26 -Block Sparse Backward - TFLOPS: 3865.68 -``` - -## Acknowledgement - -We learned or reuse code from FlexAtteniton, NATEN, and ThunderKittens. diff --git a/csrc/attn/video_sparse_attn/config_vsa.py b/csrc/attn/video_sparse_attn/config_vsa.py deleted file mode 100644 index 40e725426..000000000 --- a/csrc/attn/video_sparse_attn/config_vsa.py +++ /dev/null @@ -1,15 +0,0 @@ -### ADD TO THIS TO REGISTER NEW KERNELS -sources = { - 'block_sparse': { - 'source_files': { - 'h100': 'vsa/block_sparse_h100.cu' - } - } -} - -### WHICH KERNELS DO WE WANT TO BUILD? -# (oftentimes during development work you don't need to redefine them all.) -kernels = ['block_sparse'] - -### WHICH GPU TARGET DO WE WANT TO BUILD FOR? -target = 'h100' diff --git a/csrc/attn/video_sparse_attn/setup.py b/csrc/attn/video_sparse_attn/setup.py deleted file mode 100644 index f59303f3c..000000000 --- a/csrc/attn/video_sparse_attn/setup.py +++ /dev/null @@ -1,81 +0,0 @@ -import os -import subprocess - -from config_vsa import kernels, sources, target -from setuptools import find_packages, setup -from torch.utils.cpp_extension import BuildExtension, CUDAExtension - -target = target.lower() - -# Package metadata -PACKAGE_NAME = "vsa" -VERSION = "0.0.3" -AUTHOR = "Hao AI Lab" -DESCRIPTION = "Video Sparse Attention Kernel Used in FastVideo" -URL = "https://github.com/hao-ai-lab/FastVideo/tree/main/csrc/attn/video_sparse_attn" - -# Set environment variables -tk_root = os.getenv('THUNDERKITTENS_ROOT', os.path.abspath(os.path.join(os.getcwd(), 'tk/'))) -python_include = subprocess.check_output(['python', '-c', - "import sysconfig; print(sysconfig.get_path('include'))"]).decode().strip() -torch_include = 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() -print('vsa root:', tk_root) -print('Python include:', python_include) -print('Torch include directories:', torch_include) - -# CUDA flags -cuda_flags = [ - '-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' -] + torch_include.split() -cpp_flags = ['-std=c++20', '-O3'] - -if target == 'h100': - cuda_flags.append('-DKITTENS_HOPPER') - cuda_flags.append('-arch=sm_90a') -else: - raise ValueError(f'Target {target} not supported') - -source_files = ['vsa.cpp'] -for k in kernels: - if target not in sources[k]['source_files']: - raise KeyError(f'Target {target} not found in source files for kernel {k}') - if isinstance(sources[k]['source_files'][target], list): - source_files.extend(sources[k]['source_files'][target]) - else: - source_files.append(sources[k]['source_files'][target]) - cpp_flags.append(f'-DTK_COMPILE_{k.replace(" ", "_").upper()}') - - -ext_modules = [ - CUDAExtension('vsa_cuda', - sources=source_files, - extra_compile_args={ - 'cxx': cpp_flags, - 'nvcc': cuda_flags - }, - libraries=['cuda']) -] - - - -setup(name=PACKAGE_NAME, - version=VERSION, - author=AUTHOR, - description=DESCRIPTION, - url=URL, - packages=find_packages(), - ext_modules=ext_modules, - cmdclass={'build_ext': BuildExtension}, - classifiers=[ - "Programming Language :: Python :: 3", - "Environment :: GPU :: NVIDIA CUDA :: 12", - "License :: OSI Approved :: Apache Software License", - ], - python_requires='>=3.10', - install_requires=["torch>=2.5.0"]) diff --git a/csrc/attn/video_sparse_attn/tk b/csrc/attn/video_sparse_attn/tk deleted file mode 160000 index 6c27e28c8..000000000 --- a/csrc/attn/video_sparse_attn/tk +++ /dev/null @@ -1 +0,0 @@ -Subproject commit 6c27e28c8115d1839d9eeeb530913c184a75fc87 diff --git a/csrc/attn/video_sparse_attn/vsa.cpp b/csrc/attn/video_sparse_attn/vsa.cpp deleted file mode 100644 index 8661103a0..000000000 --- a/csrc/attn/video_sparse_attn/vsa.cpp +++ /dev/null @@ -1,27 +0,0 @@ -#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/attn/video_sparse_attn/vsa/__init__.py b/csrc/attn/video_sparse_attn/vsa/__init__.py deleted file mode 100644 index 930e981c7..000000000 --- a/csrc/attn/video_sparse_attn/vsa/__init__.py +++ /dev/null @@ -1,80 +0,0 @@ -import torch -from typing import Tuple -block_sparse_attn=None -import torch -major, minor = torch.cuda.get_device_capability(0) -if major == 9 and minor == 0:# check if H100 - from vsa_cuda import block_sparse_fwd, block_sparse_bwd - from vsa.block_sparse_wrapper import block_sparse_attn_SM90 - block_sparse_attn = block_sparse_attn_SM90 -else: - from vsa.block_sparse_wrapper import block_sparse_attn_triton - block_sparse_fwd = None - block_sparse_bwd = None - block_sparse_attn = block_sparse_attn_triton - -BLOCK_M = 64 -BLOCK_N = 64 - - -def torch_attention(q, k, v) -> Tuple[torch.Tensor, torch.Tensor]: - QK = torch.matmul(q, k.transpose(-2, -1)) - QK /= (q.size(-1)**0.5) - - # Causal mask removed since causal is always false - - QK = torch.nn.functional.softmax(QK, dim=-1) - output = torch.matmul(QK, v) - return output, QK - - -def video_sparse_attn(q, k, v, variable_block_sizes, topk, block_size, compress_attn_weight=None): - """ - q: [batch_size, num_heads, seq_len, head_dim] - k: [batch_size, num_heads, seq_len, head_dim] - v: [batch_size, num_heads, seq_len, head_dim] - topk: int - block_size: int or tuple of 3 ints - video_shape: tuple of (T, H, W) - compress_attn_weight: [batch_size, num_heads, seq_len, head_dim] - select_attn_weight: [batch_size, num_heads, seq_len, head_dim] - NOTE: We assume q, k, v is zero padded!! - V1 of sparse attention. Include compress attn and sparse attn branch, use average pooling to compress. - Assume q, k, v is flattened in this way: [batch_size, num_heads, T//block_size[0], H//block_size[1], W//block_size[2], block_size[0], block_size[1], block_size[2]] - """ - - if isinstance(block_size, int): - block_size = (block_size, block_size, block_size) - - block_elements = block_size[0] * block_size[1] * block_size[2] - assert block_elements == 64 - assert q.shape[2] % block_elements == 0 - batch_size, num_heads, seq_len, head_dim = q.shape - # compress attn - q_compress = (q.view(batch_size, num_heads, seq_len // block_elements, - block_elements, head_dim).float().sum(dim=3) / variable_block_sizes.view(1, 1, -1, 1)).to(q.dtype) - k_compress = (k.view(batch_size, num_heads, seq_len // block_elements, - block_elements, head_dim).float().sum(dim=3) / variable_block_sizes.view(1, 1, -1, 1)).to(k.dtype) - v_compress = (v.view(batch_size, num_heads, seq_len // block_elements, - block_elements, head_dim).float().sum(dim=3) / variable_block_sizes.view(1, 1, -1, 1)).to(v.dtype) - - output_compress, block_attn_score = torch_attention(q_compress, k_compress, - v_compress) - - output_compress = output_compress.view(batch_size, num_heads, - seq_len // block_elements, 1, - head_dim) - output_compress = output_compress.repeat(1, 1, 1, block_elements, - 1).view(batch_size, num_heads, - seq_len, head_dim) - - topK_indices = torch.topk(block_attn_score, topk, dim=-1).indices - block_mask = torch.zeros_like(block_attn_score, dtype=torch.bool).scatter_(-1, topK_indices, True) - output_select, _ = block_sparse_attn(q, k, v, block_mask, variable_block_sizes) - - if compress_attn_weight is not None: - final_output = output_compress * compress_attn_weight + output_select - else: - final_output = output_compress + output_select - return final_output - diff --git a/csrc/attn/video_sparse_attn/vsa/block_sparse_attn_triton.py b/csrc/attn/video_sparse_attn/vsa/block_sparse_attn_triton.py deleted file mode 100644 index 3e2c89b6e..000000000 --- a/csrc/attn/video_sparse_attn/vsa/block_sparse_attn_triton.py +++ /dev/null @@ -1,450 +0,0 @@ -""" -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 pytest -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 - - - for blk_idx in range(kv_blocks*2): - kv_idx = tl.load(kv_ptr + blk_idx//2).to(tl.int32) - block_size = tl.load(variable_block_sizes + kv_idx) - (blk_idx % 2) * step_n - block_sparse_offset = (kv_idx*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/attn/video_sparse_attn/vsa/block_sparse_h100.cu b/csrc/attn/video_sparse_attn/vsa/block_sparse_h100.cu deleted file mode 100644 index f33d2719e..000000000 --- a/csrc/attn/video_sparse_attn/vsa/block_sparse_h100.cu +++ /dev/null @@ -1,1190 +0,0 @@ -// # 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 - - // ! the following two line seems unnecessary. - // tma::store_async_wait(); // ensure qg is 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(); -} - -#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 kv_block_size -) -{ - CHECK_INPUT(q); - CHECK_INPUT(k); - CHECK_INPUT(v); - - // q shape: (batch, qo_heads, q_seq_len, head_dim) - // k shape: (batch, kv_heads, kv_seq_len, head_dim) - // v shape: (batch, kv_heads, kv_seq_len, head_dim) - // q2k_block_sparse_index shape: (batch, qo_heads, num_q_blocks, max_kv_blocks_per_q) - // q2k_block_sparse_num shape: (batch, qo_heads, num_q_blocks) - // kv_block_size shape: (num_kv_blocks) This does not need other dimensions because across all batch/heads the padding is the same. - - auto batch = q.size(0); - auto q_seq_len = q.size(2); - auto kv_seq_len = k.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 = q2k_block_sparse_index.size(2); - auto num_kv_blocks = kv_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 * BLOCK_M == q_seq_len, "This kernel supports variable q block size, but it assumes the input sequence is properly padded."); - TORCH_CHECK(num_kv_blocks * BLOCK_M == kv_seq_len, "This kernel supports variable kv block size, but it assumes the input sequence is properly padded."); - // 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(v.size(2) == kv_seq_len, "V sequence length dimension - idx 2 - must match K inputs"); - TORCH_CHECK(q2k_block_sparse_num.size(2) == num_q_blocks, "q2k_block_sparse_num idx 2 - must match num_q_blocks"); - - - 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(q_seq_len), - static_cast(head_dim)}, v.options()); - - torch::Tensor l_vec = torch::empty({static_cast(batch), - static_cast(qo_heads), - static_cast(q_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 == 64) { - using q_tile = st_bf::qo_height, fwd_attend_ker_tile_dims<64>::tile_width>; - using k_tile = st_bf::kv_height, fwd_attend_ker_tile_dims<64>::tile_width>; - using v_tile = st_bf::kv_height, fwd_attend_ker_tile_dims<64>::tile_width>; - using l_col_vec = col_vec::qo_height, fwd_attend_ker_tile_dims<64>::tile_width>>; - using o_tile = st_bf::qo_height, fwd_attend_ker_tile_dims<64>::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<64>; - - q_global qg_arg{d_q, static_cast(batch), static_cast(qo_heads), static_cast(q_seq_len), 64U}; - k_global kg_arg{d_k, static_cast(batch), static_cast(kv_heads), static_cast(kv_seq_len), 64U}; - v_global vg_arg{d_v, static_cast(batch), static_cast(kv_heads), static_cast(kv_seq_len), 64U}; - l_global lg_arg{d_l, static_cast(batch), static_cast(qo_heads), 1U, static_cast(q_seq_len)}; - o_global og_arg{d_o, static_cast(batch), static_cast(qo_heads), static_cast(q_seq_len), 64U}; - - globals g{ - qg_arg, - kg_arg, - vg_arg, - lg_arg, - og_arg, - static_cast(q_seq_len), - static_cast(hr), - static_cast(max_kv_blocks_per_q), - reinterpret_cast(q2k_block_sparse_index.data_ptr()), - reinterpret_cast(q2k_block_sparse_num.data_ptr()), - reinterpret_cast(kv_block_size.data_ptr()) - }; - - constexpr int mem_size = 54000; - - dim3 grid(q_seq_len/(BLOCK_M), qo_heads, batch); - - cudaFuncSetAttribute( - fwd_attend_ker<64>, - cudaFuncAttributeMaxDynamicSharedMemorySize, - mem_size - ); - - fwd_attend_ker<64><<>>(g); - - CHECK_CUDA_ERROR(cudaGetLastError()); - // cudaStreamSynchronize(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(q_seq_len), 128U}; - k_global kg_arg{d_k, static_cast(batch), static_cast(kv_heads), static_cast(kv_seq_len), 128U}; - v_global vg_arg{d_v, static_cast(batch), static_cast(kv_heads), static_cast(kv_seq_len), 128U}; - l_global lg_arg{d_l, static_cast(batch), static_cast(qo_heads), 1U, static_cast(q_seq_len)}; - o_global og_arg{d_o, static_cast(batch), static_cast(qo_heads), static_cast(q_seq_len), 128U}; - - globals g{ - qg_arg, - kg_arg, - vg_arg, - lg_arg, - og_arg, - static_cast(q_seq_len), - static_cast(hr), - static_cast(max_kv_blocks_per_q), - reinterpret_cast(q2k_block_sparse_index.data_ptr()), - reinterpret_cast(q2k_block_sparse_num.data_ptr()), - reinterpret_cast(kv_block_size.data_ptr()) - }; - - constexpr int mem_size = 54000; - - dim3 grid(q_seq_len/(BLOCK_M), qo_heads, batch); - - cudaFuncSetAttribute( - fwd_attend_ker<128>, - cudaFuncAttributeMaxDynamicSharedMemorySize, - mem_size - ); - - fwd_attend_ker<128><<>>(g); - - CHECK_CUDA_ERROR(cudaGetLastError()); - // cudaStreamSynchronize(stream); - } - - 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 kv_block_size) -{ - CHECK_INPUT(q); - CHECK_INPUT(k); - CHECK_INPUT(v); - CHECK_INPUT(l_vec); - CHECK_INPUT(o); - CHECK_INPUT(og); - - // q: [batch, qo_heads, q_seq_len, head_dim] - // k: [batch, kv_heads, kv_seq_len, head_dim] - // v: [batch, kv_heads, kv_seq_len, head_dim] - // o: [batch, qo_heads, q_seq_len, head_dim] - // l_vec: [batch, qo_heads, q_seq_len, 1] - // og: [batch, qo_heads, q_seq_len, head_dim] - // k2q_block_sparse_index: [batch, kv_heads, num_kv_blocks, max_num_q_blocks] - // k2q_block_sparse_num: [batch, kv_heads, num_kv_blocks] - // kv_block_size: [num_kv_blocks] - - auto batch = q.size(0); - auto q_seq_len = q.size(2); - auto kv_seq_len = k.size(2); - auto head_dim = q.size(3); - auto max_q_blocks_per_kv = k2q_block_sparse_index.size(3); - auto num_kv_blocks = kv_block_size.size(0); - TORCH_CHECK(k2q_block_sparse_index.size(2) == num_kv_blocks, "k2q_block_sparse_index.size(2) must match num_kv_blocks (kv_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(v.size(2) == kv_seq_len, "V sequence length dimension - idx 2 - must match K sequence length"); - TORCH_CHECK(l_vec.size(2) == q_seq_len, "L sequence length dimension - idx 2 - must match Q sequence length"); - TORCH_CHECK(o.size(2) == q_seq_len, "O sequence length dimension - idx 2 - must match Q sequence length"); - TORCH_CHECK(og.size(2) == q_seq_len, "OG sequence length dimension - idx 2 - must match Q sequence length"); - TORCH_CHECK(k2q_block_sparse_index.size(2) == num_kv_blocks, "k2q_block_sparse_index idx 2 - must match num_kv_blocks (kv_block_size.size(0))"); - TORCH_CHECK(k2q_block_sparse_num.size(2) == num_kv_blocks, "k2q_block_sparse_num idx 2 - must match num_kv_blocks (kv_block_size.size(0))"); - - 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(q_seq_len), - static_cast(head_dim)}, l_vec.options()); - torch::Tensor kg = torch::zeros({static_cast(batch), - static_cast(kv_heads), - static_cast(kv_seq_len), - static_cast(head_dim)}, l_vec.options()); - torch::Tensor vg = torch::zeros({static_cast(batch), - static_cast(kv_heads), - static_cast(kv_seq_len), - static_cast(head_dim)}, l_vec.options()); - - torch::Tensor d_vec = torch::empty({static_cast(batch), - static_cast(qo_heads), - static_cast(q_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(q_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(q_seq_len), 64U}; - o_global prep_o_arg {d_o, static_cast(batch), static_cast(qo_heads), static_cast(q_seq_len), 64U}; - d_global prep_d_arg {d_d, static_cast(batch), static_cast(qo_heads), 1U, static_cast(q_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(q_seq_len), 64U}; - bwd_k_global bwd_k_arg {d_k, static_cast(batch), static_cast(kv_heads), static_cast(kv_seq_len), 64U}; - bwd_v_global bwd_v_arg {d_v, static_cast(batch), static_cast(kv_heads), static_cast(kv_seq_len), 64U}; - bwd_og_global bwd_og_arg{d_og, static_cast(batch), static_cast(qo_heads), static_cast(q_seq_len), 64U}; - bwd_qg_global bwd_qg_arg{d_qg, static_cast(batch), static_cast(qo_heads), static_cast(q_seq_len), 64U}; - bwd_kg_global bwd_kg_arg{d_kg, static_cast(batch), static_cast(kv_heads), static_cast(kv_seq_len), 64U}; - bwd_vg_global bwd_vg_arg{d_vg, static_cast(batch), static_cast(kv_heads), static_cast(kv_seq_len), 64U}; - bwd_l_global bwd_l_arg {d_l, static_cast(batch), static_cast(qo_heads), 1U, static_cast(q_seq_len)}; - bwd_d_global bwd_d_arg {d_d, static_cast(batch), static_cast(qo_heads), 1U, static_cast(q_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(kv_seq_len), // N is not used in the kernel - 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(kv_block_size.data_ptr())}; - - dim3 grid_bwd_2(kv_seq_len/BLOCK_N, 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(q_seq_len), 128U}; - o_global prep_o_arg {d_o, static_cast(batch), static_cast(qo_heads), static_cast(q_seq_len), 128U}; - d_global prep_d_arg {d_d, static_cast(batch), static_cast(qo_heads), 1U, static_cast(q_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(q_seq_len), 128U}; - bwd_k_global bwd_k_arg {d_k, static_cast(batch), static_cast(kv_heads), static_cast(kv_seq_len), 128U}; - bwd_v_global bwd_v_arg {d_v, static_cast(batch), static_cast(kv_heads), static_cast(kv_seq_len), 128U}; - bwd_og_global bwd_og_arg{d_og, static_cast(batch), static_cast(qo_heads), static_cast(q_seq_len), 128U}; - bwd_qg_global bwd_qg_arg{d_qg, static_cast(batch), static_cast(qo_heads), static_cast(q_seq_len), 128U}; - bwd_kg_global bwd_kg_arg{d_kg, static_cast(batch), static_cast(kv_heads), static_cast(kv_seq_len), 128U}; - bwd_vg_global bwd_vg_arg{d_vg, static_cast(batch), static_cast(kv_heads), static_cast(kv_seq_len), 128U}; - bwd_l_global bwd_l_arg {d_l, static_cast(batch), static_cast(qo_heads), 1U, static_cast(q_seq_len)}; - bwd_d_global bwd_d_arg {d_d, static_cast(batch), static_cast(qo_heads), 1U, static_cast(q_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(kv_seq_len), // N is not used in the kernel - 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(kv_block_size.data_ptr())}; - - dim3 grid_bwd_2(kv_seq_len/BLOCK_N, 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/attn/video_sparse_attn/vsa/block_sparse_wrapper.py b/csrc/attn/video_sparse_attn/vsa/block_sparse_wrapper.py deleted file mode 100644 index bb588c44c..000000000 --- a/csrc/attn/video_sparse_attn/vsa/block_sparse_wrapper.py +++ /dev/null @@ -1,185 +0,0 @@ -import torch -try: - from vsa_cuda import block_sparse_fwd, block_sparse_bwd -except ImportError: - block_sparse_fwd = None - block_sparse_bwd = None -from vsa.block_sparse_attn_triton import triton_block_sparse_attn_forward, triton_block_sparse_attn_backward -assert torch.__version__ >= "2.4.0", "VSA requires PyTorch 2.4.0 or higher" -from vsa.index import map_to_index -from typing import Tuple, Optional - - - -@torch.library.custom_op("vsa::block_sparse_attn_triton", mutates_args=(), device_types="cuda") -def block_sparse_attn_triton( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - block_map: torch.Tensor, - variable_block_sizes: torch.Tensor, -) -> Tuple[torch.Tensor, torch.Tensor]: - q = q.contiguous() - k = k.contiguous() - v = v.contiguous() - block_map = block_map.int() - q2k_block_sparse_index, q2k_block_sparse_num = map_to_index(block_map) - o, M = triton_block_sparse_attn_forward(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, variable_block_sizes) - return o, M - - -@torch.library.register_fake("vsa::block_sparse_attn_triton") -def _block_sparse_attn_triton_fake( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - block_map: torch.Tensor, - variable_block_sizes: torch.Tensor, -) -> Tuple[torch.Tensor, torch.Tensor]: - q = q.contiguous() - k = k.contiguous() - v = v.contiguous() - o = torch.empty_like(q) - M = torch.empty((q.shape[0], q.shape[1], q.shape[2]), device=q.device, dtype=torch.float32) - return o, M - - -@torch.library.custom_op("vsa::block_sparse_attn_backward_triton", mutates_args=(), device_types="cuda") -def block_sparse_attn_backward_triton( - grad_output_padded: torch.Tensor, - q_padded: torch.Tensor, - k_padded: torch.Tensor, - v_padded: torch.Tensor, - o_padded: torch.Tensor, - M: torch.Tensor, - block_map: torch.Tensor, - variable_block_sizes: torch.Tensor, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - grad_output_padded = grad_output_padded.contiguous() - q2k_block_sparse_index, q2k_block_sparse_num = map_to_index(block_map) - k2q_block_sparse_index, k2q_block_sparse_num = map_to_index(block_map.transpose(-1, -2)) - dq, dk, dv = triton_block_sparse_attn_backward(grad_output_padded, q_padded, k_padded, v_padded, o_padded, M, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num, variable_block_sizes) - return dq, dk, dv - -@torch.library.register_fake("vsa::block_sparse_attn_backward_triton") -def _block_sparse_attn_backward_triton_fake( - grad_output_padded: torch.Tensor, - q_padded: torch.Tensor, - k_padded: torch.Tensor, - v_padded: torch.Tensor, - o_padded: torch.Tensor, - M: torch.Tensor, - block_map: torch.Tensor, - variable_block_sizes: torch.Tensor, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - grad_output_padded = grad_output_padded.contiguous() - dq = torch.empty_like(grad_output_padded) - dk = torch.empty_like(grad_output_padded) - dv = torch.empty_like(grad_output_padded) - return dq, dk, dv - - -def backward_triton(ctx, grad_output1, grad_output2): - q_padded, k_padded, v_padded, o_padded, M, block_map, variable_block_sizes = ctx.saved_tensors - dq, dk, dv = block_sparse_attn_backward_triton(grad_output1, q_padded, k_padded, v_padded, o_padded, M, block_map, variable_block_sizes) - return dq, dk, dv, None, None - - -def setup_context_triton(ctx, inputs, output): - q_padded, k_padded, v_padded, block_map, variable_block_sizes = inputs - o_padded, M = output - ctx.save_for_backward(q_padded, k_padded, v_padded, o_padded, M, block_map, variable_block_sizes) - -block_sparse_attn_triton.register_autograd(backward_triton, setup_context=setup_context_triton) - - -major, minor = torch.cuda.get_device_capability(0) - -if major == 9 and minor == 0:# check if H100 - @torch.library.custom_op("vsa::block_sparse_attn_SM90", mutates_args=(), device_types="cuda") - def block_sparse_attn_SM90( - q_padded: torch.Tensor, - k_padded: torch.Tensor, - v_padded: torch.Tensor, - block_map: torch.Tensor, - variable_block_sizes: torch.Tensor, - )-> Tuple[torch.Tensor, torch.Tensor]: - q_padded = q_padded.contiguous() - k_padded = k_padded.contiguous() - v_padded = v_padded.contiguous() - q2k_block_sparse_index, q2k_block_sparse_num = map_to_index(block_map) - variable_block_sizes = variable_block_sizes.int() - o_padded, lse_padded = block_sparse_fwd(q_padded, k_padded, v_padded, q2k_block_sparse_index, q2k_block_sparse_num, variable_block_sizes) - return o_padded, lse_padded - - - - - @torch.library.register_fake("vsa::block_sparse_attn_SM90") - def _block_sparse_attn_SM90_fake( - q_padded: torch.Tensor, - k_padded: torch.Tensor, - v_padded: torch.Tensor, - block_map: torch.Tensor, - variable_block_sizes: torch.Tensor, - ) -> Tuple[torch.Tensor, torch.Tensor]: - q_padded, k_padded, v_padded = [x.contiguous() for x in (q_padded, k_padded, v_padded)] - B, H, S, D = q_padded.shape - o_padded = torch.empty_like(q_padded) - lse_padded = torch.empty((B, H, S, 1), device=q_padded.device, dtype=torch.float32) - return o_padded, lse_padded - - - @torch.library.custom_op("vsa::block_sparse_attn_backward_SM90", mutates_args=(), device_types="cuda") - def block_sparse_attn_backward_SM90( - grad_output_padded: torch.Tensor, - q_padded: torch.Tensor, - k_padded: torch.Tensor, - v_padded: torch.Tensor, - o_padded: torch.Tensor, - lse_padded: torch.Tensor, - block_map: torch.Tensor, - variable_block_sizes: torch.Tensor, - )-> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - grad_output_padded = grad_output_padded.contiguous() - k2q_block_sparse_index, k2q_block_sparse_num = map_to_index(block_map.transpose(-1, -2)) - grad_q_padded, grad_k_padded, grad_v_padded = block_sparse_bwd( - q_padded, k_padded, v_padded, o_padded, lse_padded, grad_output_padded, k2q_block_sparse_index, k2q_block_sparse_num, variable_block_sizes - ) - grad_q_padded = grad_q_padded.to(grad_output_padded.dtype) - grad_k_padded = grad_k_padded.to(grad_output_padded.dtype) - grad_v_padded = grad_v_padded.to(grad_output_padded.dtype) - return grad_q_padded, grad_k_padded, grad_v_padded - - @torch.library.register_fake("vsa::block_sparse_attn_backward_SM90") - def _block_sparse_attn_backward_SM90_fake( - grad_output_padded: torch.Tensor, - q_padded: torch.Tensor, - k_padded: torch.Tensor, - v_padded: torch.Tensor, - o_padded: torch.Tensor, - lse_padded: torch.Tensor, - block_map: torch.Tensor, - variable_block_sizes: torch.Tensor, - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - torch._check(grad_output_padded.dtype == torch.bfloat16) - torch._check(lse_padded.dtype == torch.float32) - grad_output_padded = grad_output_padded.contiguous() - dq = torch.empty_like(grad_output_padded) - dk = torch.empty_like(grad_output_padded) - dv = torch.empty_like(grad_output_padded) - return dq, dk, dv - - - def backward_SM90(ctx, grad_output1, grad_output2): - q_padded, k_padded, v_padded, o_padded, lse_padded, block_map, variable_block_sizes= ctx.saved_tensors - dq, dk, dv = block_sparse_attn_backward_SM90(grad_output1, q_padded, k_padded, v_padded, o_padded, lse_padded, block_map, variable_block_sizes) - return dq, dk, dv, None, None - - def setup_context_SM90(ctx, inputs, output): - q_padded, k_padded, v_padded, block_map, variable_block_sizes = inputs - o_padded, lse_padded = output - ctx.save_for_backward(q_padded, k_padded, v_padded, o_padded, lse_padded, block_map, variable_block_sizes) - - - block_sparse_attn_SM90.register_autograd(backward_SM90, setup_context=setup_context_SM90) \ No newline at end of file diff --git a/csrc/attn/vmoba_attn/README.md b/csrc/attn/vmoba_attn/README.md deleted file mode 100644 index 578c4974e..000000000 --- a/csrc/attn/vmoba_attn/README.md +++ /dev/null @@ -1,32 +0,0 @@ -# Attention Kernel Used in FastVideo - -## VMoBA: Mixture-of-Block Attention for Video Diffusion Models (VMoBA) - -### Installation -Please ensure that you have installed FlashAttention version **2.7.1 or higher**, as some interfaces have changed in recent releases. - -### Usage - -You can use `moba_attn_varlen` in the following ways: - -**Install from source:** -```bash -python setup.py install -``` - -**Import after installation:** -```python -from vmoba import moba_attn_varlen -``` - -**Or import directly from the project root:** -```python -from csrc.attn.vmoba_attn.vmoba import moba_attn_varlen -``` - -### Verify if you have successfully installed - -```bash -python csrc/attn/vmoba_attn/vmoba/vmoba.py -``` - diff --git a/csrc/attn/vmoba_attn/setup.py b/csrc/attn/vmoba_attn/setup.py deleted file mode 100644 index 498f6ab5e..000000000 --- a/csrc/attn/vmoba_attn/setup.py +++ /dev/null @@ -1,26 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 - -from setuptools import find_packages, setup - -PACKAGE_NAME = "vmoba" -VERSION = "0.0.0" -AUTHOR = "JianzongWu" -DESCRIPTION = "VMoBA: Mixture-of-Block Attention for Video Diffusion Models" -URL = "https://github.com/KwaiVGI/VMoBA" - -setup( - name=PACKAGE_NAME, - version=VERSION, - author=AUTHOR, - description=DESCRIPTION, - url=URL, - packages=find_packages(), - classifiers=[ - "Programming Language :: Python :: 3", - "License :: OSI Approved :: Apache Software License", - ], - python_requires='>=3.12', - install_requires=[ - "flash-attn >= 2.7.1", - ] -) diff --git a/csrc/attn/vmoba_attn/tests/test_vmoba_attn.py b/csrc/attn/vmoba_attn/tests/test_vmoba_attn.py deleted file mode 100644 index 7e21f93df..000000000 --- a/csrc/attn/vmoba_attn/tests/test_vmoba_attn.py +++ /dev/null @@ -1,97 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 - -import torch -import pytest -import random -from csrc.attn.vmoba_attn.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.float32, 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/csrc/attn/vmoba_attn/vmoba/__init__.py b/csrc/attn/vmoba_attn/vmoba/__init__.py deleted file mode 100644 index 8119387c3..000000000 --- a/csrc/attn/vmoba_attn/vmoba/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -from .vmoba import moba_attn_varlen, process_moba_input, process_moba_output diff --git a/csrc/fastvideo_kernel/.gitignore b/csrc/fastvideo_kernel/.gitignore deleted file mode 100644 index 317ddafc4..000000000 --- a/csrc/fastvideo_kernel/.gitignore +++ /dev/null @@ -1,7 +0,0 @@ -build/ -dist/ -*.egg-info/ -__pycache__/ -*.so -*.pyc -.ipynb_checkpoints/ diff --git a/csrc/fastvideo_kernel/MANIFEST.in b/csrc/fastvideo_kernel/MANIFEST.in deleted file mode 100644 index 6d2e724a1..000000000 --- a/csrc/fastvideo_kernel/MANIFEST.in +++ /dev/null @@ -1,6 +0,0 @@ -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/csrc/st_attn.cpp b/csrc/fastvideo_kernel/csrc/st_attn.cpp deleted file mode 100644 index 41a63c17f..000000000 --- a/csrc/fastvideo_kernel/csrc/st_attn.cpp +++ /dev/null @@ -1,23 +0,0 @@ -#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/vsa.cpp b/csrc/fastvideo_kernel/csrc/vsa.cpp deleted file mode 100644 index 8661103a0..000000000 --- a/csrc/fastvideo_kernel/csrc/vsa.cpp +++ /dev/null @@ -1,27 +0,0 @@ -#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 deleted file mode 100644 index 2d2f72748..000000000 --- a/csrc/fastvideo_kernel/pyproject.toml +++ /dev/null @@ -1,27 +0,0 @@ -[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 deleted file mode 100644 index 455306704..000000000 --- a/csrc/fastvideo_kernel/setup.py +++ /dev/null @@ -1,132 +0,0 @@ -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/triton_kernels/index.py b/csrc/fastvideo_kernel/src/fastvideo_kernel/triton_kernels/index.py deleted file mode 100644 index 6225f15b4..000000000 --- a/csrc/fastvideo_kernel/src/fastvideo_kernel/triton_kernels/index.py +++ /dev/null @@ -1,152 +0,0 @@ - -## 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 deleted file mode 100644 index 6bac66ac6..000000000 --- a/csrc/fastvideo_kernel/src/fastvideo_kernel/vmoba.py +++ /dev/null @@ -1,868 +0,0 @@ -# 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/docker/Dockerfile.python3.10 b/docker/Dockerfile.python3.10 index a39a416a6..98a1e0bd7 100644 --- a/docker/Dockerfile.python3.10 +++ b/docker/Dockerfile.python3.10 @@ -55,18 +55,12 @@ RUN source $HOME/.local/bin/env && \ echo 'source /opt/venv/bin/activate' >> /root/.bashrc && \ echo 'if [ -n "$ZSH_VERSION" ] && [ -f ~/.zshrc ]; then . ~/.zshrc; elif [ -f ~/.bashrc ]; then . ~/.bashrc; fi' > /root/.profile -# Install STA (Sliding Tile Attention) +# Install FastVideo Unified Kernel RUN source $HOME/.local/bin/env && \ source /opt/venv/bin/activate && \ - cd csrc/attn/sliding_tile_attn && \ + cd fastvideo-kernel && \ git submodule update --init --recursive && \ - python setup.py install + ./build.sh -# Install VSA -RUN source $HOME/.local/bin/env && \ - source /opt/venv/bin/activate && \ - cd csrc/attn/video_sparse_attn && \ - git submodule update --init --recursive && \ - python setup.py install EXPOSE 22 diff --git a/docker/Dockerfile.python3.11 b/docker/Dockerfile.python3.11 index 769b4a1fa..fe3639a60 100644 --- a/docker/Dockerfile.python3.11 +++ b/docker/Dockerfile.python3.11 @@ -55,18 +55,12 @@ RUN source $HOME/.local/bin/env && \ echo 'source /opt/venv/bin/activate' >> /root/.bashrc && \ echo 'if [ -n "$ZSH_VERSION" ] && [ -f ~/.zshrc ]; then . ~/.zshrc; elif [ -f ~/.bashrc ]; then . ~/.bashrc; fi' > /root/.profile -# Install STA (Sliding Tile Attention) +# Install FastVideo Unified Kernel RUN source $HOME/.local/bin/env && \ source /opt/venv/bin/activate && \ - cd csrc/attn/sliding_tile_attn && \ + cd fastvideo-kernel && \ git submodule update --init --recursive && \ - python setup.py install + ./build.sh -# Install VSA -RUN source $HOME/.local/bin/env && \ - source /opt/venv/bin/activate && \ - cd csrc/attn/video_sparse_attn && \ - git submodule update --init --recursive && \ - python setup.py install EXPOSE 22 diff --git a/docker/Dockerfile.python3.12 b/docker/Dockerfile.python3.12 index f1ebe560a..64c0103c8 100644 --- a/docker/Dockerfile.python3.12 +++ b/docker/Dockerfile.python3.12 @@ -55,18 +55,11 @@ RUN source $HOME/.local/bin/env && \ echo 'source /opt/venv/bin/activate' >> /root/.bashrc && \ echo 'if [ -n "$ZSH_VERSION" ] && [ -f ~/.zshrc ]; then . ~/.zshrc; elif [ -f ~/.bashrc ]; then . ~/.bashrc; fi' > /root/.profile -# Install STA (Sliding Tile Attention) +# Install FastVideo Unified Kernel RUN source $HOME/.local/bin/env && \ source /opt/venv/bin/activate && \ - cd csrc/attn/sliding_tile_attn && \ + cd fastvideo-kernel && \ git submodule update --init --recursive && \ - python setup.py install - -# Install VSA -RUN source $HOME/.local/bin/env && \ - source /opt/venv/bin/activate && \ - cd csrc/attn/video_sparse_attn && \ - git submodule update --init --recursive && \ - python setup.py install + ./build.sh EXPOSE 22 diff --git a/docker/Dockerfile.python3.12.cuda12.9.1 b/docker/Dockerfile.python3.12.cuda12.9.1 index 5ecaf95d0..ce90131d8 100644 --- a/docker/Dockerfile.python3.12.cuda12.9.1 +++ b/docker/Dockerfile.python3.12.cuda12.9.1 @@ -55,18 +55,12 @@ RUN source $HOME/.local/bin/env && \ echo 'source /opt/venv/bin/activate' >> /root/.bashrc && \ echo 'if [ -n "$ZSH_VERSION" ] && [ -f ~/.zshrc ]; then . ~/.zshrc; elif [ -f ~/.bashrc ]; then . ~/.bashrc; fi' > /root/.profile -# Install STA (Sliding Tile Attention) +# Install FastVideo Unified Kernel RUN source $HOME/.local/bin/env && \ source /opt/venv/bin/activate && \ - cd csrc/attn/sliding_tile_attn && \ + cd fastvideo-kernel && \ git submodule update --init --recursive && \ - python setup.py install + ./build.sh -# Install VSA -RUN source $HOME/.local/bin/env && \ - source /opt/venv/bin/activate && \ - cd csrc/attn/video_sparse_attn && \ - git submodule update --init --recursive && \ - python setup.py install EXPOSE 22 diff --git a/docker/rocm7_1.Dockerfile b/docker/rocm7_1.Dockerfile index 96cfc5cd2..0ceba9651 100644 --- a/docker/rocm7_1.Dockerfile +++ b/docker/rocm7_1.Dockerfile @@ -48,11 +48,11 @@ RUN source $HOME/.local/bin/env && \ echo 'source /opt/venv/bin/activate' >> /root/.bashrc && \ echo 'if [ -n "$ZSH_VERSION" ] && [ -f ~/.zshrc ]; then . ~/.zshrc; elif [ -f ~/.bashrc ]; then . ~/.bashrc; fi' > /root/.profile -# Install STA (Sliding Tile Attention) +# Install FastVideo Unified Kernel RUN source $HOME/.local/bin/env && \ source /opt/venv/bin/activate && \ - cd csrc/attn/sliding_tile_attn && \ + cd fastvideo-kernel && \ git submodule update --init --recursive && \ - SKIP_SM90_EXT=1 pip install --no-build-isolation -e . + ./build.sh EXPOSE 22 diff --git a/docs/sliding_tile_attention/installation.md b/docs/sliding_tile_attention/installation.md index 25d183ab2..7da08aced 100644 --- a/docs/sliding_tile_attention/installation.md +++ b/docs/sliding_tile_attention/installation.md @@ -28,24 +28,24 @@ export PATH=${CUDA_HOME}/bin:${PATH} export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH ``` -Install STA: +Install FastVideo Kernel: ```bash -cd csrc/attn/sliding_tile_attn/ +cd fastvideo-kernel/ git submodule update --init --recursive -python setup.py install +./build.sh ``` # 🧪 Test ```bash -python csrc/attn/tests/test_sta.py +pytest fastvideo-kernel/tests/test_correctness.py ``` # 📋 Usage ```python -from st_attn import sliding_tile_attention +from fastvideo_kernel import sliding_tile_attention # assuming video size (T, H, W) = (30, 48, 80), text tokens = 256 with padding. # q, k, v: [batch_size, num_heads, seq_length, head_dim], seq_length = T*H*W + 256 # a tile is a cube of size (6, 8, 8) diff --git a/docs/video_sparse_attention/installation.md b/docs/video_sparse_attention/installation.md index 6ca95bd6d..74b5ca7e6 100644 --- a/docs/video_sparse_attention/installation.md +++ b/docs/video_sparse_attention/installation.md @@ -29,24 +29,24 @@ export PATH=${CUDA_HOME}/bin:${PATH} export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH ``` -Install VSA: +Install FastVideo Kernel: ```bash -cd csrc/attn/video_sparse_attn/ +cd fastvideo-kernel/ git submodule update --init --recursive -python setup.py install +./build.sh ``` # 🧪 Test ```bash -python csrc/attn/tests/test_vsa.py +pytest fastvideo-kernel/tests/test_correctness.py ``` # 📋 Usage ```python -from vsa import video_sparse_attn +from fastvideo_kernel import video_sparse_attn # q, k, v: [batch_size, num_heads, seq_len, head_dim] # variable_block_sizes: [num_blocks] - number of valid tokens in each block diff --git a/fastvideo-kernel/CMakeLists.txt b/fastvideo-kernel/CMakeLists.txt new file mode 100644 index 000000000..ba6e5d3b8 --- /dev/null +++ b/fastvideo-kernel/CMakeLists.txt @@ -0,0 +1,114 @@ +cmake_minimum_required(VERSION 3.26 FATAL_ERROR) +project(fastvideo-kernel LANGUAGES CXX CUDA) + +# Import common utils if needed, but we keep it simple for now + +# Find Python and Torch +find_package(Python COMPONENTS Interpreter Development.Module REQUIRED) + +# Robustly find Torch include paths using Python +execute_process( + COMMAND "${Python_EXECUTABLE}" -c "import torch; from torch.utils.cpp_extension import include_paths; print(';'.join(include_paths()))" + OUTPUT_VARIABLE TORCH_INCLUDE_PATHS + OUTPUT_STRIP_TRAILING_WHITESPACE +) +list(APPEND TORCH_INCLUDE_DIRS ${TORCH_INCLUDE_PATHS}) + +# Find Torch package (still useful for libraries) +find_package(Torch REQUIRED) + +# Include directories +include_directories( + ${CMAKE_SOURCE_DIR}/include + ${CMAKE_SOURCE_DIR}/include/tk/include + ${CMAKE_SOURCE_DIR}/include/tk/prototype + ${CMAKE_SOURCE_DIR}/csrc + ${TORCH_INCLUDE_DIRS} +) + +# CUDA architecture flags +# If not defined, try to detect from PyTorch +if(NOT DEFINED TORCH_CUDA_ARCH_LIST) + execute_process( + COMMAND "${Python_EXECUTABLE}" -c "import torch; print(';'.join([x.replace('sm_', '') for x in torch.cuda.get_arch_list()]))" + OUTPUT_VARIABLE TORCH_CUDA_ARCH_LIST + OUTPUT_STRIP_TRAILING_WHITESPACE + ) +endif() + +if(NOT TORCH_CUDA_ARCH_LIST) + # Fallback default if detection failed or no GPU + set(TORCH_CUDA_ARCH_LIST "9.0a") +endif() + +message(STATUS "TORCH_CUDA_ARCH_LIST: ${TORCH_CUDA_ARCH_LIST}") + +# Check if we are targeting Hopper (9.0a) for ThunderKittens +string(FIND "${TORCH_CUDA_ARCH_LIST}" "9.0a" HAS_HOPPER) +if(HAS_HOPPER EQUAL -1) + message(STATUS "TORCH_CUDA_ARCH_LIST (${TORCH_CUDA_ARCH_LIST}) does not contain 9.0a. ThunderKittens kernels will be disabled.") + set(ENABLE_TK_KERNELS OFF) +else() + message(STATUS "Enabling ThunderKittens kernels (Hopper supported).") + set(ENABLE_TK_KERNELS ON) +endif() + +# Always try to build the extension if CUDA is available, but conditionally add sources/flags +set(BUILD_CXX_KERNELS ON) + +# Compiler flags +set(CUDA_FLAGS + "-DNDEBUG" + "-O3" + "-std=c++20" + "--use_fast_math" + "--expt-extended-lambda" + "--expt-relaxed-constexpr" + "-Xcompiler=-fno-strict-aliasing" + "-Xcompiler=-fPIC" + "-DTORCH_COMPILE" + "-Xnvlink=--verbose" + "-Xptxas=--verbose" + "-Xptxas=--warn-on-spills" +) + +# Only add -arch=sm_90a and TK flags if we are enabling TK kernels +if(ENABLE_TK_KERNELS) + list(APPEND CUDA_FLAGS "-arch=sm_90a") + list(APPEND CUDA_FLAGS "-DKITTENS_HOPPER") +endif() + +if(BUILD_CXX_KERNELS) + # Source files + set(EXTENSION_SOURCES csrc/common_extension.cpp) + + # Conditionally add TK kernels + if(ENABLE_TK_KERNELS) + list(APPEND EXTENSION_SOURCES + csrc/attention/st_attn_h100.cu + csrc/attention/block_sparse_h100.cu + ) + endif() + + # Combined FastVideo Extension + # Using name 'fastvideo_kernel_ops' to distinguish from the python package namespace + Python_add_library(fastvideo_kernel_ops MODULE USE_SABI ${SKBUILD_SABI_VERSION} WITH_SOABI + ${EXTENSION_SOURCES} + ) + + # Build compile definitions list + set(COMPILE_DEFS TORCH_EXTENSION_NAME=fastvideo_kernel_ops) + if(ENABLE_TK_KERNELS) + list(APPEND COMPILE_DEFS TK_COMPILE_ST_ATTN TK_COMPILE_BLOCK_SPARSE) + endif() + + target_compile_definitions(fastvideo_kernel_ops PRIVATE ${COMPILE_DEFS}) + + target_compile_options(fastvideo_kernel_ops PRIVATE + $<$:${CUDA_FLAGS}> + ) + + # We install it to fastvideo_kernel/_C so we can load it to register the ops + install(TARGETS fastvideo_kernel_ops LIBRARY DESTINATION fastvideo_kernel/_C) +endif() + diff --git a/csrc/fastvideo_kernel/LICENSE b/fastvideo-kernel/LICENSE similarity index 100% rename from csrc/fastvideo_kernel/LICENSE rename to fastvideo-kernel/LICENSE diff --git a/fastvideo-kernel/MANIFEST.in b/fastvideo-kernel/MANIFEST.in new file mode 100644 index 000000000..695242784 --- /dev/null +++ b/fastvideo-kernel/MANIFEST.in @@ -0,0 +1,6 @@ +include LICENSE +include README.md +include pyproject.toml +recursive-include python/fastvideo_kernel *.py +recursive-include csrc *.cu *.cuh *.cpp *.h +recursive-include include/tk *.cu *.cuh *.cpp *.h *.src diff --git a/csrc/fastvideo_kernel/README.md b/fastvideo-kernel/README.md similarity index 96% rename from csrc/fastvideo_kernel/README.md rename to fastvideo-kernel/README.md index 72412db3f..95bcd3016 100644 --- a/csrc/fastvideo_kernel/README.md +++ b/fastvideo-kernel/README.md @@ -6,7 +6,7 @@ CUDA kernels for FastVideo video generation. ```bash git submodule update --init --recursive -cd csrc/fastvideo_kernel +cd fastvideo-kernel pip install . ``` diff --git a/fastvideo-kernel/build.sh b/fastvideo-kernel/build.sh new file mode 100755 index 000000000..b480cee78 --- /dev/null +++ b/fastvideo-kernel/build.sh @@ -0,0 +1,21 @@ +#!/bin/bash +set -ex + +# Simple build script wrapping uv/pip + +echo "Building fastvideo-kernel..." + +# Ensure submodules are initialized if needed (tk) +# git submodule update --init --recursive + +# Install build dependencies +pip install scikit-build-core cmake ninja + +# Set TORCH_CUDA_ARCH_LIST only if not already set +if [ -z "$TORCH_CUDA_ARCH_LIST" ]; then + export TORCH_CUDA_ARCH_LIST="9.0a" +fi + +# Build and install +# Use -v for verbose output +pip install . -v --no-build-isolation diff --git a/csrc/fastvideo_kernel/csrc/block_sparse_h100.cu b/fastvideo-kernel/csrc/attention/block_sparse_h100.cu similarity index 100% rename from csrc/fastvideo_kernel/csrc/block_sparse_h100.cu rename to fastvideo-kernel/csrc/attention/block_sparse_h100.cu diff --git a/csrc/fastvideo_kernel/csrc/st_attn_h100.cu b/fastvideo-kernel/csrc/attention/st_attn_h100.cu similarity index 100% rename from csrc/fastvideo_kernel/csrc/st_attn_h100.cu rename to fastvideo-kernel/csrc/attention/st_attn_h100.cu diff --git a/fastvideo-kernel/csrc/common_extension.cpp b/fastvideo-kernel/csrc/common_extension.cpp new file mode 100644 index 000000000..78f07007a --- /dev/null +++ b/fastvideo-kernel/csrc/common_extension.cpp @@ -0,0 +1,36 @@ +#include +#include +#include + +// Forward declarations +#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 + +#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() = "FastVideo CUDA Kernels"; + +#ifdef TK_COMPILE_ST_ATTN + m.def("sta_fwd", torch::wrap_pybind_function(sta_forward), "sliding tile attention (Hopper)"); +#endif + +#ifdef TK_COMPILE_BLOCK_SPARSE + m.def("block_sparse_fwd", torch::wrap_pybind_function(block_sparse_attention_forward), "block sparse attention forward (Hopper)"); + m.def("block_sparse_bwd", torch::wrap_pybind_function(block_sparse_attention_backward), "block sparse attention backward (Hopper)"); +#endif +} diff --git a/csrc/attn/sliding_tile_attn/tk b/fastvideo-kernel/include/tk similarity index 100% rename from csrc/attn/sliding_tile_attn/tk rename to fastvideo-kernel/include/tk diff --git a/fastvideo-kernel/pyproject.toml b/fastvideo-kernel/pyproject.toml new file mode 100644 index 000000000..fc043c8f2 --- /dev/null +++ b/fastvideo-kernel/pyproject.toml @@ -0,0 +1,35 @@ +[build-system] +requires = [ + "scikit-build-core>=0.10", + "torch>=2.5.0", + "setuptools>=61.0.0", + "wheel" +] +build-backend = "scikit_build_core.build" + +[project] +name = "fastvideo-kernel" +version = "0.2.0" +description = "Unified CUDA kernels for FastVideo" +readme = "README.md" +requires-python = ">=3.10" +license = { file = "LICENSE" } +authors = [ + { name = "Hao AI Lab", email = "contact@haoailab.com" } +] +classifiers = [ + "Programming Language :: Python :: 3", + "Environment :: GPU :: NVIDIA CUDA", +] +dependencies = [ + "torch>=2.5.0", + "triton>=2.0.0", +] + +[project.urls] +"Homepage" = "https://github.com/hao-ai-lab/FastVideo" + +[tool.scikit-build] +cmake.build-type = "Release" +minimum-version = "build-system.requires" +wheel.packages = ["python"] diff --git a/csrc/fastvideo_kernel/src/fastvideo_kernel/__init__.py b/fastvideo-kernel/python/fastvideo_kernel/__init__.py similarity index 100% rename from csrc/fastvideo_kernel/src/fastvideo_kernel/__init__.py rename to fastvideo-kernel/python/fastvideo_kernel/__init__.py diff --git a/csrc/fastvideo_kernel/src/fastvideo_kernel/ops.py b/fastvideo-kernel/python/fastvideo_kernel/ops.py similarity index 65% rename from csrc/fastvideo_kernel/src/fastvideo_kernel/ops.py rename to fastvideo-kernel/python/fastvideo_kernel/ops.py index b5feeacdc..5dfe31e92 100644 --- a/csrc/fastvideo_kernel/src/fastvideo_kernel/ops.py +++ b/fastvideo-kernel/python/fastvideo_kernel/ops.py @@ -1,20 +1,20 @@ import math import torch from .triton_kernels.block_sparse_attn_triton import triton_block_sparse_attn_forward +from .triton_kernels.st_attn_triton import sliding_tile_attention_triton from .triton_kernels.index import map_to_index +# Try to load the C++ extension try: - from fastvideo_kernel._C.st_attn import sta_fwd + from fastvideo_kernel._C import fastvideo_kernel_ops + sta_fwd = getattr(fastvideo_kernel_ops, "sta_fwd", None) + block_sparse_fwd = getattr(fastvideo_kernel_ops, "block_sparse_fwd", None) + block_sparse_bwd = getattr(fastvideo_kernel_ops, "block_sparse_bwd", None) 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, @@ -24,12 +24,15 @@ def sliding_tile_attention( has_text: bool = True, seq_shape: str = "30x48x80", ) -> torch.Tensor: + # Check if the specific op is available if sta_fwd is None: - raise RuntimeError("STA kernel not compiled. Requires H100 and ThunderKittens at build time.") - + return sliding_tile_attention_triton( + q, k, v, window_size, text_length, has_text, seq_shape + ) + 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 @@ -37,22 +40,20 @@ def sliding_tile_attention( 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], + 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] @@ -67,37 +68,45 @@ def video_sparse_attn( ) -> 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) + + 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) - + 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) + mask = torch.zeros_like(scores, + dtype=torch.bool).scatter_(-1, topk_idx, True) + + idx, num = map_to_index(mask) 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()) + out_s = block_sparse_fwd( + q, k, v, idx, num, variable_block_sizes.int() + )[0] # block_sparse_fwd returns vector else: - idx, num = map_to_index(mask) - out_s, _ = triton_block_sparse_attn_forward(q, k, v, idx, num, variable_block_sizes) - + 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/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/block_sparse_attn_triton.py similarity index 55% rename from csrc/fastvideo_kernel/src/fastvideo_kernel/triton_kernels/block_sparse_attn_triton.py rename to fastvideo-kernel/python/fastvideo_kernel/triton_kernels/block_sparse_attn_triton.py index 373e2f2ae..b828cc2d1 100644 --- a/csrc/fastvideo_kernel/src/fastvideo_kernel/triton_kernels/block_sparse_attn_triton.py +++ b/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/block_sparse_attn_triton.py @@ -8,7 +8,6 @@ This is a Triton implementation of the Flash Attention v2 algorithm from Tri Dao Credits: OpenAI kernel team """ - import torch import triton import triton.language as tl @@ -17,7 +16,6 @@ import triton.language as tl 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. @@ -29,65 +27,92 @@ configs = [ 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): +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_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 + 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) + 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)) + 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)) + 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) + 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)) + 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) @@ -127,23 +152,30 @@ def _attn_fwd_sparse(Q, K, V, sm_scale, # 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 # - ): +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) + 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) @@ -151,19 +183,32 @@ def _attn_bwd_preprocess(O, DO, # # 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): +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) @@ -172,21 +217,20 @@ def _attn_bwd_dkdv(dk, dv, # # 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 + 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 + 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 + + 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) @@ -212,21 +256,32 @@ def _attn_bwd_dkdv(dk, dv, # 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): +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) @@ -237,27 +292,27 @@ def _attn_bwd_dq(dq, q, K, V, # # 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_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 + 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 + + 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) + 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]) @@ -269,23 +324,37 @@ def _attn_bwd_dq(dq, q, K, V, # 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): +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) @@ -319,20 +388,32 @@ def _attn_bwd(Q, K, V, sm_scale, # 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, # + dk, + dv, # + Q, + k, + v, + sm_scale, # DO, # - M, D, # - k2q_index, k2q_num, max_q_blks, + 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 # + 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 @@ -357,54 +438,88 @@ def _attn_bwd(Q, K, V, sm_scale, # 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 # - ) + 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): +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]}" + assert q2k_num.shape[ + -1] == T // 64, 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 - ) + _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): + +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) @@ -420,30 +535,49 @@ def triton_block_sparse_attn_backward(do, q, k, v, o, M, q2k_index, q2k_num, k2q pre_grid = (N_CTX // PRE_BLOCK, BATCH * N_HEAD) delta = torch.empty_like(M) _attn_bwd_preprocess[pre_grid]( - o, do, # + o, + do, # delta, # - BATCH, N_HEAD, N_CTX, # - BLOCK_M=PRE_BLOCK, HEAD_DIM=D # + 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, + 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, # + 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/attn/video_sparse_attn/vsa/index.py b/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/index.py similarity index 99% rename from csrc/attn/video_sparse_attn/vsa/index.py rename to fastvideo-kernel/python/fastvideo_kernel/triton_kernels/index.py index 6225f15b4..d84d01650 100644 --- a/csrc/attn/video_sparse_attn/vsa/index.py +++ b/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/index.py @@ -1,9 +1,9 @@ - ## 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, @@ -26,6 +26,7 @@ def topk_index_to_map_kernel( 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, @@ -59,6 +60,7 @@ def map_to_index_kernel( 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): @@ -106,6 +108,7 @@ def topk_index_to_map(index: torch.Tensor, return block_map + def map_to_index(block_map: torch.Tensor): """ Convert a block map to indices and counts. diff --git a/csrc/attn/sliding_tile_attn/st_attn/st_attn_triton.py b/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/st_attn_triton.py similarity index 100% rename from csrc/attn/sliding_tile_attn/st_attn/st_attn_triton.py rename to fastvideo-kernel/python/fastvideo_kernel/triton_kernels/st_attn_triton.py diff --git a/csrc/attn/vmoba_attn/vmoba/vmoba.py b/fastvideo-kernel/python/fastvideo_kernel/vmoba.py similarity index 52% rename from csrc/attn/vmoba_attn/vmoba/vmoba.py rename to fastvideo-kernel/python/fastvideo_kernel/vmoba.py index 6bac66ac6..8f95c5455 100644 --- a/csrc/attn/vmoba_attn/vmoba/vmoba.py +++ b/fastvideo-kernel/python/fastvideo_kernel/vmoba.py @@ -3,19 +3,21 @@ 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`.") + 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 @@ -37,23 +39,25 @@ def calc_chunks(cu_seqlen, moba_chunk_size): ) 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 = 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 + 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 = 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 - ) + 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 @@ -61,12 +65,11 @@ def calc_chunks(cu_seqlen, moba_chunk_size): # --- 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: + +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). @@ -75,45 +78,53 @@ def _select_threshold_query_head( 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 + 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) + 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 + 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) + 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 + 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) + 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) + 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) + 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) + 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) + 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) @@ -125,12 +136,9 @@ def _select_threshold_query_head( 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: +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). @@ -140,53 +148,64 @@ def _select_threshold_block( 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 + 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_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) + 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) + 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) + 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,) + 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,) + 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 + 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) + 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) + 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,) + 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,) + 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) + 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 = 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) + 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) @@ -194,12 +213,9 @@ def _select_threshold_block( 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: +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. @@ -209,66 +225,79 @@ def _select_threshold_overall( 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 + 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_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) + 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) + 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) + 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 + 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 + 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 + 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 - + 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 + 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,) + 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 - + 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) - ) + 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 = 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) + 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) @@ -276,12 +305,10 @@ def _select_threshold_overall( 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: +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. """ @@ -307,7 +334,8 @@ def _select_threshold_head_global( 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 = 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" @@ -321,25 +349,28 @@ def _select_threshold_head_global( .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) + 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) + 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,) + 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) + 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) + 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) + 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) @@ -348,6 +379,7 @@ def _select_threshold_head_global( class MixedAttention(torch.autograd.Function): + @staticmethod def forward( ctx, @@ -365,7 +397,7 @@ class MixedAttention(torch.autograd.Function): ): ctx.max_seqlen = max_seqlen ctx.moba_chunk_size = moba_chunk_size - ctx.softmax_scale = softmax_scale = q.shape[-1] ** (-0.5) + ctx.softmax_scale = softmax_scale = q.shape[-1]**(-0.5) # Non-causal self-attention branch # return out, softmax_lse, S_dmask, rng_state @@ -398,26 +430,25 @@ class MixedAttention(torch.autograd.Function): 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 = 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) - ) + 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_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 @@ -426,17 +457,15 @@ class MixedAttention(torch.autograd.Function): 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) - ) + 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) + 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, @@ -500,46 +529,40 @@ class MixedAttention(torch.autograd.Function): alibi_slopes=None, deterministic=True, window_size_left=-1, - window_size_right=-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) - ) + 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) - ) + 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 - ) + _ = _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 @@ -552,7 +575,7 @@ def moba_attn_varlen( max_seqlen: int, moba_chunk_size: int, moba_topk: int, - select_mode: str = 'threshold', # "topk" or "threshold" + select_mode: str = 'threshold', # "topk" or "threshold" simsum_threshold: float = 0.25, threshold_type: str = 'query_head', ) -> torch.Tensor: @@ -573,8 +596,7 @@ def moba_attn_varlen( # Compute chunk boundaries. cu_chunk, filtered_chunk_indices, num_filtered_chunk, chunk_to_batch = calc_chunks( - cu_seqlens, moba_chunk_size - ) + cu_seqlens, moba_chunk_size) self_attn_cu_seqlen = cu_chunk @@ -582,20 +604,25 @@ def moba_attn_varlen( 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] + 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) + 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) + 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] + 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) @@ -604,14 +631,19 @@ def moba_attn_varlen( # 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) + 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) + 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) + 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() @@ -632,7 +664,11 @@ def moba_attn_varlen( 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_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 @@ -643,56 +679,85 @@ def moba_attn_varlen( 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) + 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) + 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_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'))) + 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) + _, 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'." - ) + "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) + 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) + 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) + 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) + 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'.") + 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'.") - + 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 - ) + 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) + 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() @@ -710,19 +775,23 @@ def moba_attn_varlen( # 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) + 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) + 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_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)], + [ + torch.zeros(1, device=experts_tensor.device, dtype=torch.int32), + valid_expert_lengths.cumsum(dim=0) + ], dim=0, ).to(torch.int32) @@ -762,22 +831,48 @@ def process_moba_input( 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]) + 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)}" + 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]) + 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]) + 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)}") - + 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 @@ -789,9 +884,23 @@ def process_moba_output( 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]) + 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]) + 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 @@ -803,11 +912,18 @@ def generate_data(batch_size, seqlen, num_head, head_dim, dtype): 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') + 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') + 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) ...") @@ -816,53 +932,104 @@ def generate_data(batch_size, seqlen, num_head, head_dim, dtype): 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'): +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}") + 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) + 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) + 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) + 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) + 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 + python -u fastvideo-kernel/python/fastvideo_kernel/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') + 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/src/fastvideo_kernel/_C/__init__.py b/fastvideo-kernel/tests/__init__.py similarity index 100% rename from csrc/fastvideo_kernel/src/fastvideo_kernel/_C/__init__.py rename to fastvideo-kernel/tests/__init__.py diff --git a/csrc/fastvideo_kernel/tests/support_flex_sta.py b/fastvideo-kernel/tests/support_flex_sta.py similarity index 100% rename from csrc/fastvideo_kernel/tests/support_flex_sta.py rename to fastvideo-kernel/tests/support_flex_sta.py diff --git a/csrc/fastvideo_kernel/tests/test_correctness.py b/fastvideo-kernel/tests/test_correctness.py similarity index 100% rename from csrc/fastvideo_kernel/tests/test_correctness.py rename to fastvideo-kernel/tests/test_correctness.py diff --git a/csrc/fastvideo_kernel/tests/test_vmoba_correctness.py b/fastvideo-kernel/tests/test_vmoba_correctness.py similarity index 98% rename from csrc/fastvideo_kernel/tests/test_vmoba_correctness.py rename to fastvideo-kernel/tests/test_vmoba_correctness.py index c5d1a7f72..ff5db6091 100644 --- a/csrc/fastvideo_kernel/tests/test_vmoba_correctness.py +++ b/fastvideo-kernel/tests/test_vmoba_correctness.py @@ -3,7 +3,7 @@ import torch import pytest import random -from fastvideo_kernel.vmoba import moba_attn_varlen +from fastvideo_kernel import moba_attn_varlen def generate_test_data(batch_size, total_seqlen, num_heads, head_dim, dtype, device="cuda"): """ diff --git a/fastvideo/attention/backends/sdpa.py b/fastvideo/attention/backends/sdpa.py index 00a0cee60..f7960f5e0 100644 --- a/fastvideo/attention/backends/sdpa.py +++ b/fastvideo/attention/backends/sdpa.py @@ -81,7 +81,9 @@ class SDPAImpl(AttentionImpl): query = query.transpose(1, 2) key = key.transpose(1, 2) value = value.transpose(1, 2) - attn_mask = attn_metadata.attn_mask if attn_metadata is not None else None + + if attn_metadata is not None: + attn_mask = getattr(attn_metadata, "attn_mask", None) attn_kwargs = { "attn_mask": attn_mask, "dropout_p": self.dropout, diff --git a/fastvideo/platforms/rocm.py b/fastvideo/platforms/rocm.py index 2b1445d71..3d71a2a25 100644 --- a/fastvideo/platforms/rocm.py +++ b/fastvideo/platforms/rocm.py @@ -80,7 +80,7 @@ class RocmPlatform(Platform): elif 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) diff --git a/fastvideo/tests/modal/pr_test.py b/fastvideo/tests/modal/pr_test.py index 140b920cf..8afc6a0d6 100644 --- a/fastvideo/tests/modal/pr_test.py +++ b/fastvideo/tests/modal/pr_test.py @@ -53,6 +53,7 @@ def run_test(pytest_command: str): git clone {git_repo} /FastVideo && cd /FastVideo && {checkout_command} && + uv pip install -e fastvideo-kernel && uv pip install -e .[test] && {pytest_command} """ @@ -103,15 +104,16 @@ def run_inference_tests_STA(): @app.function(gpu="H100:1", image=image, timeout=900) def run_precision_tests_STA(): - run_test("python csrc/attn/tests/test_sta.py") + run_test("pytest fastvideo-kernel/tests/test_correctness.py") @app.function(gpu="H100:1", image=image, timeout=900) def run_precision_tests_VSA(): - run_test("python csrc/attn/tests/test_vsa.py") + # VSA correctness is covered by the same file now + run_test("pytest fastvideo-kernel/tests/test_correctness.py") @app.function(gpu="L40S:1", image=image, timeout=900) def run_precision_tests_vmoba(): - run_test("pytest csrc/attn/vmoba_attn/tests/test_vmoba_attn.py") + run_test("pytest fastvideo-kernel/tests/test_vmoba_correctness.py") @app.function(gpu="L40S:1", image=image, timeout=900) def run_inference_tests_vmoba():