diff --git a/fastvideo/attention/backends/sageattn/api.py b/fastvideo/attention/backends/sageattn/api.py index 28b11f4f1..def4f50fe 100644 --- a/fastvideo/attention/backends/sageattn/api.py +++ b/fastvideo/attention/backends/sageattn/api.py @@ -24,8 +24,8 @@ import fp4quant_cuda # Centralized block size configuration for sageattn_blackwell kernels # These should match the values in fastvideo/attention/backends/sageattn/blackwell/block_config.h -BLOCK_M = 128 # Block size for M dimension (query sequence length) -BLOCK_N = 128 # Block size for N dimension (key/value sequence length) +BLOCK_M = 64 # Block size for M dimension (query sequence length) +BLOCK_N = 64 # Block size for N dimension (key/value sequence length) @triton.jit diff --git a/fastvideo/attention/backends/sageattn/blackwell/block_config.h b/fastvideo/attention/backends/sageattn/blackwell/block_config.h index e4ecfda77..631e605ae 100644 --- a/fastvideo/attention/backends/sageattn/blackwell/block_config.h +++ b/fastvideo/attention/backends/sageattn/blackwell/block_config.h @@ -20,9 +20,9 @@ // Block sizes for M and N dimensions namespace flash { // Block size for M dimension (query sequence length) - static constexpr int BLOCK_M = 128; + static constexpr int BLOCK_M = 64; // Block size for N dimension (key/value sequence length) - static constexpr int BLOCK_N = 128; + static constexpr int BLOCK_N = 64; }