diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index b13f19345..99b54b6b7 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -14,13 +14,9 @@ on: - ".github/workflows/pr-test.yml" - "pyproject.toml" - "docker/Dockerfile.python3.12" + - "csrc/**" workflow_dispatch: inputs: - custom_image: - description: "Custom image from this repository (default: fastvideo-dev:py3.12-latest)" - required: false - default: "fastvideo-dev:py3.12-latest" - type: string run_encoder_test: description: "Run encoder-test" required: false @@ -56,6 +52,16 @@ on: required: false default: false type: boolean + run_precision_test_STA: + description: "Run precision-test-STA" + required: false + default: false + type: boolean + run_precision_test_VSA: + description: "Run precision-test-VSA" + required: false + default: false + type: boolean run_nightly_test: description: "Run nightly-test" required: false @@ -65,6 +71,7 @@ on: env: PYTHONUNBUFFERED: "1" + concurrency: group: pr-test-${{ github.ref }} cancel-in-progress: true @@ -84,44 +91,69 @@ jobs: training-test: ${{ steps.filter.outputs.training-test }} training-test-VSA: ${{ steps.filter.outputs.training-test-VSA }} inference-test-STA: ${{ steps.filter.outputs.inference-test-STA }} + precision-test-STA: ${{ steps.filter.outputs.precision-test-STA }} + precision-test-VSA: ${{ steps.filter.outputs.precision-test-VSA }} steps: - uses: actions/checkout@v4 - uses: dorny/paths-filter@v3 id: filter with: filters: | + # Define reusable path patterns + common-paths: &common-paths + - 'pyproject.toml' + - 'docker/Dockerfile.python3.12' + sta-kernel-paths: &sta-kernel-paths + - 'csrc/attn/st_attn/**' + - 'csrc/attn/setup_sta.py' + - 'csrc/attn/config_sta.py' + - 'csrc/attn/st_attn.cpp' + vsa-kernel-paths: &vsa-kernel-paths + - 'csrc/attn/vsa/**' + - 'csrc/attn/tk/**' + - 'csrc/attn/setup_vsa.py' + - 'csrc/attn/config_vsa.py' + - 'csrc/attn/vsa.cpp' + vsa-paths: &vsa-paths + - 'fastvideo/v1/**' + - *common-paths + - *vsa-kernel-paths + + # Actual tests encoder-test: - 'fastvideo/v1/models/encoders/**' - 'fastvideo/v1/models/loaders/**' - 'fastvideo/v1/tests/encoders/**' - - 'pyproject.toml' - - 'docker/Dockerfile.python3.12' + - *common-paths vae-test: - 'fastvideo/v1/models/vaes/**' - 'fastvideo/v1/models/loaders/**' - 'fastvideo/v1/tests/vaes/**' - - 'pyproject.toml' - - 'docker/Dockerfile.python3.12' + - *common-paths transformer-test: - 'fastvideo/v1/models/dits/**' - 'fastvideo/v1/models/loaders/**' - 'fastvideo/v1/tests/transformers/**' - 'fastvideo/v1/layers/**' - 'fastvideo/v1/attention/**' - - 'pyproject.toml' - - 'docker/Dockerfile.python3.12' + - *common-paths training-test: - 'fastvideo/v1/**' - - 'pyproject.toml' - - 'docker/Dockerfile.python3.12' + - *common-paths training-test-VSA: - 'fastvideo/v1/**' - - 'pyproject.toml' - - 'docker/Dockerfile.python3.12' + - *common-paths + - *vsa-kernel-paths inference-test-STA: - 'fastvideo/v1/**' - - 'pyproject.toml' - - 'docker/Dockerfile.python3.12' + - *common-paths + - *sta-kernel-paths + precision-test-STA: + - *common-paths + - *sta-kernel-paths + precision-test-VSA: + - *common-paths + - *vsa-kernel-paths encoder-test: needs: change-filter @@ -134,7 +166,7 @@ jobs: gpu_type: "NVIDIA A40" gpu_count: 1 volume_size: 100 - image: "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:py3.12-latest' }}" + image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest" test_command: "uv pip install -e .[test] && pytest ./fastvideo/v1/tests/encoders -s" timeout_minutes: 30 secrets: @@ -152,7 +184,7 @@ jobs: gpu_type: "NVIDIA A40" gpu_count: 1 volume_size: 100 - image: "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:py3.12-latest' }}" + image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest" test_command: "uv pip install -e .[test] && pytest ./fastvideo/v1/tests/vaes -s" timeout_minutes: 30 secrets: @@ -170,7 +202,7 @@ jobs: gpu_type: "NVIDIA L40S" gpu_count: 1 volume_size: 100 - image: "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:py3.12-latest' }}" + image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest" test_command: "uv pip install -e .[test] && pytest ./fastvideo/v1/tests/transformers -s" timeout_minutes: 30 secrets: @@ -216,7 +248,7 @@ jobs: gpu_count: 4 volume_size: 100 disk_size: 100 - image: "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:py3.12-latest' }}" + image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest" test_command: "wandb login $WANDB_API_KEY && uv pip install -e .[test] && pytest ./fastvideo/v1/tests/training/Vanilla -srP" timeout_minutes: 30 secrets: @@ -236,7 +268,7 @@ jobs: gpu_count: 1 volume_size: 100 disk_size: 100 - image: "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:py3.12-latest' }}" + image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest" test_command: "wandb login $WANDB_API_KEY && uv pip install -e .[test] && pytest ./fastvideo/v1/tests/training/VSA -srP" timeout_minutes: 30 secrets: @@ -256,13 +288,51 @@ jobs: gpu_count: 1 volume_size: 100 disk_size: 100 - image: "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:py3.12-latest' }}" + image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest" test_command: "uv pip install -e .[test] && pytest ./fastvideo/v1/tests/inference/STA -srP" timeout_minutes: 30 secrets: RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }} RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }} + precision-test-STA: + needs: change-filter + if: >- + (github.event_name != 'workflow_dispatch' && github.event.pull_request.draft == false) || + (github.event_name == 'workflow_dispatch' && github.event.inputs.run_precision_test_STA == 'true') + uses: ./.github/workflows/runpod-test.yml + with: + job_id: "precision-test-STA" + gpu_type: "NVIDIA H100 NVL" + gpu_count: 1 + volume_size: 100 + disk_size: 100 + image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest" + test_command: "uv pip install -e .[test] && python csrc/attn/tests/test_sta.py" + timeout_minutes: 30 + secrets: + RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }} + RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }} + + precision-test-VSA: + needs: change-filter + if: >- + (github.event_name != 'workflow_dispatch' && github.event.pull_request.draft == false) || + (github.event_name == 'workflow_dispatch' && github.event.inputs.run_precision_test_VSA == 'true') + uses: ./.github/workflows/runpod-test.yml + with: + job_id: "precision-test-VSA" + gpu_type: "NVIDIA H100 NVL" + gpu_count: 1 + volume_size: 100 + disk_size: 100 + image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest" + test_command: "uv pip install -e .[test] && python csrc/attn/tests/test_block_sparse.py" + timeout_minutes: 30 + secrets: + RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }} + RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }} + nightly-test: if: >- (github.event_name == 'workflow_dispatch' && github.event.inputs.run_nightly_test == 'true') @@ -273,7 +343,7 @@ jobs: gpu_count: 4 volume_size: 100 disk_size: 100 - image: "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:py3.12-latest' }}" + image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest" test_command: "wandb login $WANDB_API_KEY && uv pip install -e .[test] && pytest ./fastvideo/v1/tests/nightly/test_e2e_overfit_single_sample.py -vs" timeout_minutes: 30 secrets: @@ -282,7 +352,8 @@ jobs: WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }} runpod-cleanup: - needs: [encoder-test, vae-test, transformer-test, ssim-test] # Add other jobs to this list as you create them + # Add other jobs to this list as you create them + needs: [encoder-test, vae-test, transformer-test, ssim-test, training-test, training-test-VSA, inference-test-STA, precision-test-STA, precision-test-VSA] if: ${{ always() && ((github.event_name != 'workflow_dispatch' && github.event.pull_request.draft == false) || github.event_name == 'workflow_dispatch') }} runs-on: ubuntu-latest steps: @@ -299,7 +370,7 @@ jobs: - name: Cleanup all RunPod instances env: - JOB_IDS: '["encoder-test", "vae-test", "transformer-test", "ssim-test-py3.10", "ssim-test-py3.11", "ssim-test-py3.12"]' + JOB_IDS: '["encoder-test", "vae-test", "transformer-test", "ssim-test-py3.10", "ssim-test-py3.11", "ssim-test-py3.12", "training-test", "training-test-VSA", "inference-test-STA", "precision-test-STA", "precision-test-VSA"]' RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }} GITHUB_RUN_ID: ${{ github.run_id }} run: python .github/scripts/runpod_cleanup.py diff --git a/csrc/attn/README.md b/csrc/attn/README.md index 7b0f63835..6569b3ebc 100644 --- a/csrc/attn/README.md +++ b/csrc/attn/README.md @@ -4,7 +4,7 @@ ## Installation -We test our code on Pytorch 2.5.0 and CUDA>=12.4. Currently we only have implementation on H100. +We test our code on Pytorch 2.5.0 and CUDA>=12.4. Currently we only support H100/H200, because ThunderKittens uses TMA but doesn't support Blackwell yet. First, install C++20 for ThunderKittens: ```bash sudo apt update @@ -53,8 +53,14 @@ out = sliding_tile_attention(q, k, v, window_size, 0, False) ## Test ```bash -python test/test_sta.py +python tests/test_sta.py # test STA +python tests/test_block_sparse.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. diff --git a/csrc/attn/bench/bench_sta.py b/csrc/attn/benchmarks/bench_sta.py similarity index 67% rename from csrc/attn/bench/bench_sta.py rename to csrc/attn/benchmarks/bench_sta.py index 8023f4ea6..534fa2627 100644 --- a/csrc/attn/bench/bench_sta.py +++ b/csrc/attn/benchmarks/bench_sta.py @@ -5,6 +5,7 @@ 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"): @@ -13,16 +14,16 @@ def flops(batch, seqlen, nheads, headdim, causal, mode="fwd"): return f if mode == "fwd" else (2.5 * f if mode == "bwd" else 3.5 * f) -def efficiency(flop, time): - flop = flop / 1e12 - time = time / 1e6 - return flop / time +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 in configurations: + 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}") @@ -30,38 +31,31 @@ def benchmark_attention(configurations): 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() + # 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() - 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() - # Prepare for timing forward pass - start_events_fwd = [torch.cuda.Event(enable_timing=True) for _ in range(10)] - end_events_fwd = [torch.cuda.Event(enable_timing=True) for _ in range(10)] + # # Warmup for forward pass + # for _ in range(10): + # o = sliding_tile_attention(q, k, v, [[3, 6, 10]] * 24, 0, False, dit_seq_shape) - torch.cuda.empty_cache() - torch.cuda.synchronize() + # # 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)) - # Warmup for forward pass - for _ in range(10): - o = sliding_tile_attention(q, k, v, [[3, 6, 10]] * 24, 0, False, '18x48x80') + # 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 - # 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, '18x48x80') - end_events_fwd[i].record() - - torch.cuda.synchronize() - 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 = efficiency(flops(B, N, H, D, causal, 'fwd'), time_us_fwd) + 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 in us: {time_us_fwd:.2f}") - print(f"Average efficiency for forward pass in TFLOPS: {tflops_fwd}") + print(f"Average time for forward pass (ms): {ms:.2f}") + print(f"Average TFLOPS: {tflops_fwd}") print("-" * 60) # torch.cuda.empty_cache() @@ -85,15 +79,14 @@ def benchmark_attention(configurations): # 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 = efficiency(flops(B, N, H, D, causal, 'bwd'), time_us_bwd) + # 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 in us: {time_us_bwd:.2f}") - # print(f"Average efficiency for backward pass in TFLOPS: {tflops_bwd}") - print("=" * 60) + # print(f"Average time for backward pass(ms): {ms:.2f}") + # print(f"Average TFLOPS: {tflops_bwd}") + # print("=" * 60) torch.cuda.empty_cache() - torch.cuda.synchronize() return results @@ -124,7 +117,10 @@ def plot_results(results): # Example list of configurations to test configurations = [ - (2, 24, 69120, 128, False), + (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), diff --git a/csrc/attn/bench/bench_vsa.py b/csrc/attn/benchmarks/bench_vsa.py similarity index 100% rename from csrc/attn/bench/bench_vsa.py rename to csrc/attn/benchmarks/bench_vsa.py diff --git a/csrc/attn/st_attn/st_attn_h100.cu b/csrc/attn/st_attn/st_attn_h100.cu index ef6c245db..208dc6ecc 100644 --- a/csrc/attn/st_attn/st_attn_h100.cu +++ b/csrc/attn/st_attn/st_attn_h100.cu @@ -4,9 +4,17 @@ #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; +} -#define CLAMP(value, min, max) ((value) < (min) ? (min) : ((value) > (max) ? (max) : (value))) -#define ABS(x) ((x) < 0 ? -(x) : (x)) constexpr int CONSUMER_WARPGROUPS = (3); constexpr int PRODUCER_WARPGROUPS = (1); @@ -117,16 +125,16 @@ void fwd_attend_ker(const __grid_constant__ fwd_globals g) { int qt = seq_idx / 6 / (CH * CW); int qh = (seq_idx / 6) % (CH * CW) / CW; int qw = (seq_idx / 6) % CW; - qt = CLAMP(qt, DT, CT-DT-1); - qh = CLAMP(qh, DH, CH-DH-1); - qw = CLAMP(qw, DW, CW-DW-1); + 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(qt - kt) <= DT) && (ABS(qh - kh) <= DH) && (ABS(qw - kw) <= DW); + 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)); @@ -167,15 +175,15 @@ void fwd_attend_ker(const __grid_constant__ fwd_globals g) { int qt = seq_idx / 6 / (CH * CW); int qh = (seq_idx / 6) % (CH * CW) / CW; int qw = (seq_idx / 6) % CW; - qt = CLAMP(qt, DT, CT-DT-1); - qh = CLAMP(qh, DH, CH-DH-1); - qw = CLAMP(qw, DW, CW-DW-1); - int k_t_min = CLAMP(qt-DT, 0, CT-1); - int k_t_max = CLAMP(qt+DT, 0, CT-1); - int k_h_min = CLAMP(qh-DH, 0, CH-1); - int k_h_max = CLAMP(qh+DH, 0, CH-1); - int k_w_min = CLAMP(qw-DW, 0, CW-1); - int k_w_max = CLAMP(qw+DW, 0, CW-1); + 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++) { @@ -234,7 +242,7 @@ void fwd_attend_ker(const __grid_constant__ fwd_globals g) { // the last three kv blocks are for text, we process them separately kv_iters = img_kv_blocks - 1; } else { - kv_iters = CLAMP(DT*2+1, 1, CT) * CLAMP(DH*2+1, 1, CH) * CLAMP(DW*2+1, 1, CW) * 3 - 1 ; + 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); @@ -415,8 +423,9 @@ sta_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o, float* l_ptr = reinterpret_cast(l_vec.data_ptr()); float* d_l = reinterpret_cast(l_ptr); - cudaDeviceSynchronize(); - auto stream = at::cuda::getCurrentCUDAStream().stream(); + //cudadevicesynchronize(); + const c10::cuda::OptionalCUDAGuard device_guard(q.device()); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); if (head_dim == 128) { @@ -442,8 +451,8 @@ sta_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o, globals g{qg_arg, kg_arg, vg_arg, lg_arg, og_arg, static_cast(seq_len), static_cast(text_length), static_cast(hr)}; - auto mem_size = kittens::MAX_SHARED_MEMORY; - auto threads = NUM_WORKERS * kittens::WARP_THREADS; + 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); @@ -823,10 +832,10 @@ sta_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o, } CHECK_CUDA_ERROR(cudaGetLastError()); - cudaStreamSynchronize(stream); + // cudaStreamSynchronize(stream); } return o; - cudaDeviceSynchronize(); + //cudadevicesynchronize(); } diff --git a/csrc/attn/test/flex_sta_ref.py b/csrc/attn/tests/flex_sta_ref.py similarity index 100% rename from csrc/attn/test/flex_sta_ref.py rename to csrc/attn/tests/flex_sta_ref.py diff --git a/csrc/attn/test/test_block_sparse.py b/csrc/attn/tests/test_block_sparse.py similarity index 81% rename from csrc/attn/test/test_block_sparse.py rename to csrc/attn/tests/test_block_sparse.py index 42555b6fb..ba31075d8 100644 --- a/csrc/attn/test/test_block_sparse.py +++ b/csrc/attn/tests/test_block_sparse.py @@ -7,6 +7,7 @@ from vsa import BLOCK_M, BLOCK_N import numpy as np import random +import gc def set_seed(seed: int = 42): # Python random module @@ -20,15 +21,6 @@ def set_seed(seed: int = 42): 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=4, help='Batch size') - parser.add_argument('--num_heads', type=int, default=6, help='Number of heads') - parser.add_argument('--head_dim', type=int, default=128, help='Head dimension') - parser.add_argument('--topk', type=int, default=64, help='Number of kv blocks each q block attends to') - parser.add_argument('--seq_lengths', type=int, nargs='+', default=[29120], help='Sequence lengths to benchmark') - parser.add_argument('--num_iterations', type=int, default=100, help='Number of test iterations to run') - return parser.parse_args() @torch.no_grad def precision_metric(quant_o, fa2_o): @@ -135,9 +127,7 @@ def generate_block_sparse_pattern(bs, h, num_q_blocks, num_kv_blocks, k, device= return q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num, block_sparse_mask -def main(): - args = parse_arguments() - +def main(args): set_seed(42) # Extract parameters @@ -191,23 +181,36 @@ def main(): block_mask_expanded = block_sparse_mask.unsqueeze(-1).unsqueeze(-2) # [b, h, num_q_blocks, num_kv_blocks, 1, 1] block_mask_expanded = block_mask_expanded.expand(-1, -1, -1, -1, BLOCK_M, BLOCK_N) # [b, h, num_q_blocks, num_kv_blocks, BLOCK_M, BLOCK_N] full_mask = block_mask_expanded.permute(0, 1, 2, 4, 3, 5).reshape(batch, head, seq_len, seq_len) - - q_sdpa = q.clone() - k_sdpa = k.clone() - v_sdpa = v.clone() q.requires_grad = True k.requires_grad = True v.requires_grad = True - q_sdpa.requires_grad = True - k_sdpa.requires_grad = True - v_sdpa.requires_grad = True + # testing forward o = BlockSparseAttentionFunction.apply(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num) + del q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num, block_sparse_mask, block_mask_expanded + grad_o = torch.randn_like(o) + o.backward(grad_o) + # clear memory + q_sdpa = q.detach().clone() + k_sdpa = k.detach().clone() + v_sdpa = v.detach().clone() + q_sdpa.requires_grad = True + k_sdpa.requires_grad = True + v_sdpa.requires_grad = True + q.data = torch.empty(0, device=q.device) + k.data = torch.empty(0, device=k.device) + v.data = torch.empty(0, device=v.device) + torch.cuda.empty_cache() + o_sdpa = torch.nn.functional.scaled_dot_product_attention(q_sdpa, k_sdpa, v_sdpa, attn_mask=full_mask) + sim, l1, rmse = precision_metric(o, o_sdpa) + assert sim > 0.9999, f"SSIM too low: {sim}" + assert l1 < 8e-5, f"l1 too large: {l1}" + assert rmse < 2e-5, f"RMSE too large: {rmse}" forward_metrics['sim'].append(sim) forward_metrics['l1'].append(l1) forward_metrics['rmse'].append(rmse) @@ -215,52 +218,72 @@ def main(): print(f"block_sparse_attention_fwd vs torch.nn.functional.scaled_dot_product_attention:\nsim: {sim}, l1: {l1}, rmse: {rmse}") # test backward - grad_o = torch.randn_like(o) - o.backward(grad_o) o_sdpa.backward(grad_o) sim, l1, rmse = precision_metric(q.grad, q_sdpa.grad) + # Error bounds collected on H100 + assert sim > 0.9999, f"SSIM too low: {sim}" + assert l1 < 4e-3, f"l1 too large: {l1}" + assert rmse < 3e-4, f"RMSE too large: {rmse}" grad_q_metrics['sim'].append(sim) grad_q_metrics['l1'].append(l1) grad_q_metrics['rmse'].append(rmse) print(f"block_sparse_attention_bwd vs torch.nn.functional.scaled_dot_product_attention grad_q:\nsim: {sim}, l1: {l1}, rmse: {rmse}") sim, l1, rmse = precision_metric(k.grad, k_sdpa.grad) + assert sim > 0.9999, f"SSIM too low: {sim}" + assert l1 < 4e-3, f"l1 too large: {l1}" + assert rmse < 2e-4, f"RMSE too large: {rmse}" grad_k_metrics['sim'].append(sim) grad_k_metrics['l1'].append(l1) grad_k_metrics['rmse'].append(rmse) print(f"block_sparse_attention_bwd vs torch.nn.functional.scaled_dot_product_attention grad_k:\nsim: {sim}, l1: {l1}, rmse: {rmse}") sim, l1, rmse = precision_metric(v.grad, v_sdpa.grad) + assert sim > 0.9999, f"SSIM too low: {sim}" + assert l1 < 1e-4, f"l1 too large: {l1}" + assert rmse < 2e-5, f"RMSE too large: {rmse}" grad_v_metrics['sim'].append(sim) grad_v_metrics['l1'].append(l1) grad_v_metrics['rmse'].append(rmse) print(f"block_sparse_attention_bwd vs torch.nn.functional.scaled_dot_product_attention grad_v:\nsim: {sim}, l1: {l1}, rmse: {rmse}") - + + del o, o_sdpa, grad_o, q_sdpa, k_sdpa, v_sdpa + gc.collect() + torch.cuda.empty_cache() + # Print summary statistics if multiple iterations were run if num_iterations > 1: print("\n" + "="*50) print(f"Summary Statistics (over {num_iterations} iterations):") print("\nForward metrics:") - print(f"Similarity: mean={np.mean(forward_metrics['sim']):.6f}, std={np.std(forward_metrics['sim']):.6f}") - print(f"L1 error: mean={np.mean(forward_metrics['l1']):.6f}, std={np.std(forward_metrics['l1']):.6f}") - print(f"RMSE: mean={np.mean(forward_metrics['rmse']):.6f}, std={np.std(forward_metrics['rmse']):.6f}") + print(f"Similarity: mean={np.mean(forward_metrics['sim']):.6f}, std={np.std(forward_metrics['sim']):.6f}, min={np.min(forward_metrics['sim']):.6f}") + print(f"L1 error: mean={np.mean(forward_metrics['l1']):.6f}, std={np.std(forward_metrics['l1']):.6f}, max={np.max(forward_metrics['l1']):.6f}") + print(f"RMSE: mean={np.mean(forward_metrics['rmse']):.6f}, std={np.std(forward_metrics['rmse']):.6f}, max={np.max(forward_metrics['rmse']):.6f}") print("\nGradient Q metrics:") - print(f"Similarity: mean={np.mean(grad_q_metrics['sim']):.6f}, std={np.std(grad_q_metrics['sim']):.6f}") - print(f"L1 error: mean={np.mean(grad_q_metrics['l1']):.6f}, std={np.std(grad_q_metrics['l1']):.6f}") - print(f"RMSE: mean={np.mean(grad_q_metrics['rmse']):.6f}, std={np.std(grad_q_metrics['rmse']):.6f}") + print(f"Similarity: mean={np.mean(grad_q_metrics['sim']):.6f}, std={np.std(grad_q_metrics['sim']):.6f}, min={np.min(grad_q_metrics['sim']):.6f}") + print(f"L1 error: mean={np.mean(grad_q_metrics['l1']):.6f}, std={np.std(grad_q_metrics['l1']):.6f}, max={np.max(grad_q_metrics['l1']):.6f}") + print(f"RMSE: mean={np.mean(grad_q_metrics['rmse']):.6f}, std={np.std(grad_q_metrics['rmse']):.6f}, max={np.max(grad_q_metrics['rmse']):.6f}") print("\nGradient K metrics:") - print(f"Similarity: mean={np.mean(grad_k_metrics['sim']):.6f}, std={np.std(grad_k_metrics['sim']):.6f}") - print(f"L1 error: mean={np.mean(grad_k_metrics['l1']):.6f}, std={np.std(grad_k_metrics['l1']):.6f}") - print(f"RMSE: mean={np.mean(grad_k_metrics['rmse']):.6f}, std={np.std(grad_k_metrics['rmse']):.6f}") + print(f"Similarity: mean={np.mean(grad_k_metrics['sim']):.6f}, std={np.std(grad_k_metrics['sim']):.6f}, min={np.min(grad_k_metrics['sim']):.6f}") + print(f"L1 error: mean={np.mean(grad_k_metrics['l1']):.6f}, std={np.std(grad_k_metrics['l1']):.6f}, max={np.max(grad_k_metrics['l1']):.6f}") + print(f"RMSE: mean={np.mean(grad_k_metrics['rmse']):.6f}, std={np.std(grad_k_metrics['rmse']):.6f}, max={np.max(grad_k_metrics['rmse']):.6f}") print("\nGradient V metrics:") - print(f"Similarity: mean={np.mean(grad_v_metrics['sim']):.6f}, std={np.std(grad_v_metrics['sim']):.6f}") - print(f"L1 error: mean={np.mean(grad_v_metrics['l1']):.6f}, std={np.std(grad_v_metrics['l1']):.6f}") - print(f"RMSE: mean={np.mean(grad_v_metrics['rmse']):.6f}, std={np.std(grad_v_metrics['rmse']):.6f}") + print(f"Similarity: mean={np.mean(grad_v_metrics['sim']):.6f}, std={np.std(grad_v_metrics['sim']):.6f}, min={np.min(grad_v_metrics['sim']):.6f}") + print(f"L1 error: mean={np.mean(grad_v_metrics['l1']):.6f}, std={np.std(grad_v_metrics['l1']):.6f}, max={np.max(grad_v_metrics['l1']):.6f}") + print(f"RMSE: mean={np.mean(grad_v_metrics['rmse']):.6f}, std={np.std(grad_v_metrics['rmse']):.6f}, max={np.max(grad_v_metrics['rmse']):.6f}") if __name__ == "__main__": - main() \ No newline at end of file + parser = argparse.ArgumentParser(description='Benchmark Block Sparse Attention') + parser.add_argument('--batch_size', type=int, default=4, help='Batch size') + parser.add_argument('--num_heads', type=int, default=6, help='Number of heads') + parser.add_argument('--head_dim', type=int, default=128, help='Head dimension') + parser.add_argument('--topk', type=int, default=64, help='Number of kv blocks each q block attends to') + parser.add_argument('--seq_lengths', type=int, nargs='+', default=[29120], help='Sequence lengths to benchmark') + parser.add_argument('--num_iterations', type=int, default=50, help='Number of test iterations to run') + args = parser.parse_args() + main(args) \ No newline at end of file diff --git a/csrc/attn/test/test_block_sparse_mask_to_index.py b/csrc/attn/tests/test_block_sparse_mask_to_index.py similarity index 100% rename from csrc/attn/test/test_block_sparse_mask_to_index.py rename to csrc/attn/tests/test_block_sparse_mask_to_index.py diff --git a/csrc/attn/test/test_fa2.py b/csrc/attn/tests/test_fa2.py similarity index 100% rename from csrc/attn/test/test_fa2.py rename to csrc/attn/tests/test_fa2.py diff --git a/csrc/attn/test/test_full_attn.py b/csrc/attn/tests/test_full_attn.py similarity index 100% rename from csrc/attn/test/test_full_attn.py rename to csrc/attn/tests/test_full_attn.py diff --git a/csrc/attn/test/test_sta.py b/csrc/attn/tests/test_sta.py similarity index 92% rename from csrc/attn/test/test_sta.py rename to csrc/attn/tests/test_sta.py index 2daf99859..699e56cd6 100644 --- a/csrc/attn/test/test_sta.py +++ b/csrc/attn/tests/test_sta.py @@ -81,5 +81,7 @@ 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/vsa/block_sparse_h100.cu b/csrc/attn/vsa/block_sparse_h100.cu index c3f1c04b8..2b74c7a8b 100644 --- a/csrc/attn/vsa/block_sparse_h100.cu +++ b/csrc/attn/vsa/block_sparse_h100.cu @@ -3,6 +3,8 @@ #include "kittens.cuh" #include #include +#include + using namespace kittens; namespace cg = cooperative_groups; @@ -940,8 +942,9 @@ block_sparse_attention_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v float* l_ptr = reinterpret_cast(l_vec.data_ptr()); float* d_l = reinterpret_cast(l_ptr); - cudaDeviceSynchronize(); - auto stream = at::cuda::getCurrentCUDAStream().stream(); + //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>; @@ -966,7 +969,7 @@ block_sparse_attention_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v globals g{qg_arg, kg_arg, vg_arg, lg_arg, og_arg, static_cast(seq_len), static_cast(hr), static_cast(max_kv_blocks_per_q), reinterpret_cast(q2k_block_sparse_index.data_ptr()), reinterpret_cast(q2k_block_sparse_num.data_ptr())}; - auto mem_size = 54000; + constexpr int mem_size = 54000; dim3 grid(seq_len/(64), qo_heads, batch); @@ -979,7 +982,7 @@ block_sparse_attention_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v fwd_attend_ker<64><<>>(g); CHECK_CUDA_ERROR(cudaGetLastError()); - cudaStreamSynchronize(stream); + // cudaStreamSynchronize(stream); } if (head_dim == 128) { @@ -1005,7 +1008,7 @@ block_sparse_attention_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v globals g{qg_arg, kg_arg, vg_arg, lg_arg, og_arg, static_cast(seq_len), static_cast(hr), static_cast(max_kv_blocks_per_q), reinterpret_cast(q2k_block_sparse_index.data_ptr()), reinterpret_cast(q2k_block_sparse_num.data_ptr())}; - auto mem_size = 54000; + constexpr int mem_size = 54000; dim3 grid(seq_len/(64), qo_heads, batch); @@ -1018,11 +1021,11 @@ block_sparse_attention_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v fwd_attend_ker<128><<>>(g); CHECK_CUDA_ERROR(cudaGetLastError()); - cudaStreamSynchronize(stream); + // cudaStreamSynchronize(stream); } return {o, l_vec}; - cudaDeviceSynchronize(); + //cudadevicesynchronize(); } std::vector @@ -1132,13 +1135,14 @@ block_sparse_attention_backward(torch::Tensor q, float* d_kg = reinterpret_cast(kg_ptr); float* d_vg = reinterpret_cast(vg_ptr); - auto mem_size = kittens::MAX_SHARED_MEMORY; - auto threads = 4 * kittens::WARP_THREADS; + constexpr int mem_size = kittens::MAX_SHARED_MEMORY; + int threads = 4 * kittens::WARP_THREADS; - cudaDeviceSynchronize(); - auto stream = at::cuda::getCurrentCUDAStream().stream(); + //cudadevicesynchronize(); + const c10::cuda::OptionalCUDAGuard device_guard(q.device()); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); - cudaStreamSynchronize(stream); + // cudaStreamSynchronize(stream); // TORCH_CHECK(seq_len % (4*kittens::TILE_DIM*4) == 0, "sequence length must be divisible by 256"); dim3 grid_bwd(seq_len/(4*kittens::TILE_ROW_DIM*4), qo_heads, batch); @@ -1222,7 +1226,7 @@ block_sparse_attention_backward(torch::Tensor q, dim3 grid_bwd_2(seq_len/64, qo_heads, batch); threads = 128; - cudaDeviceSynchronize(); + //cudadevicesynchronize(); { cudaFuncSetAttribute( @@ -1240,8 +1244,8 @@ block_sparse_attention_backward(torch::Tensor q, } // CHECK_CUDA_ERROR(cudaGetLastError()); - cudaStreamSynchronize(stream); - cudaDeviceSynchronize(); + // 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; @@ -1326,7 +1330,7 @@ block_sparse_attention_backward(torch::Tensor q, dim3 grid_bwd_2(seq_len/64, qo_heads, batch); threads = 128; - cudaDeviceSynchronize(); + //cudadevicesynchronize(); { cudaFuncSetAttribute( @@ -1338,10 +1342,10 @@ block_sparse_attention_backward(torch::Tensor q, bwd_attend_ker<128><<>>(bwd_global); } - cudaStreamSynchronize(stream); - cudaDeviceSynchronize(); + // cudaStreamSynchronize(stream); + //cudadevicesynchronize(); } return {qg, kg, vg}; - cudaDeviceSynchronize(); + //cudadevicesynchronize(); } \ No newline at end of file