Add benchmark for Flash Attention Decode (#363)

This PR adds the benchmark for Flash Attention Decode. It is stacked on
top of #362.

---------

Co-authored-by: Alejandro Acosta <[email protected]>

Author: muhammad-tanvir-1211

benchmarks/README.md

@@ -44,8 +44,9 @@ ninja benchmarks_gemm_sycl
 # target = intel_gpu_pvc | intel_gpu_bmg_g21
 cmake .. -GNinja -DCUTLASS_ENABLE_SYCL=ON -DDPCPP_SYCL_TARGET=$target -DCUTLASS_ENABLE_BENCHMARKS=ON -DCUTLASS_ENABLE_TESTS=ON
 
-ninja cutlass_benchmarks_flash_attention_prefill
+ninja cutlass_benchmarks_flash_attention
 ./benchmarks/flash_attention/flash_attention_prefill/cutlass_benchmarks_flash_attention_prefill_xe --config_file=../benchmarks/device/pvc/input_files/input_flash_attention_prefill.in
+./benchmarks/flash_attention/flash_attention_decode/cutlass_benchmarks_flash_attention_decode_xe --config_file=../benchmarks/device/bmg/input_files/input_sglang_flash_attention_decode_kvcache.in
 ```
 
 ## Compiling and Running Flash Attention v2 benchmarks with default configurations with Intel Xe backend
@@ -54,7 +55,7 @@ ninja cutlass_benchmarks_flash_attention_prefill
 # target = intel_gpu_pvc | intel_gpu_bmg_g21
 cmake .. -GNinja -DCUTLASS_ENABLE_SYCL=ON -DDPCPP_SYCL_TARGET=$target -DCUTLASS_ENABLE_BENCHMARKS=ON -DCUTLASS_ENABLE_TESTS=ON
 
-ninja benchmarks_flash_attention_prefill
+ninja benchmarks_flash_attention
 ```
 
 ## Compiling and Running all benchmarks with default configurations with Intel Xe backend

benchmarks/device/bmg/input_files/input_sglang_flash_attention_decode_kvcache.in

@@ -0,0 +1,31 @@
+# Flash attention decode (with kv-cache)
+
+PvcFMHADecodeBF16BF16FP32_RCR_h128_NonCausal_FixedLen --bm_name=attention_decode_kv_cache --batch=1 --seq_len_qo=1 --num_heads_q=32 --head_size_qk=128 --head_size_vo=128 --seq_len_kv=2048 --seq_len_kv_cache=1024 --num_heads_kv=8
+PvcFMHADecodeBF16BF16FP32_RCR_h128_Causal_FixedLen --bm_name=attention_decode_kv_cache --batch=1 --seq_len_qo=1 --num_heads_q=32 --head_size_qk=128 --head_size_vo=128 --seq_len_kv=2048 --seq_len_kv_cache=1024 --num_heads_kv=8
+PvcFMHADecodeBF16BF16FP32_RCR_h128_NonCausal_FixedLen --bm_name=attention_decode_kv_cache --batch=8 --seq_len_qo=1 --num_heads_q=32 --head_size_qk=128 --head_size_vo=128 --seq_len_kv=1024 --seq_len_kv_cache=1024 --num_heads_kv=8
+PvcFMHADecodeBF16BF16FP32_RCR_h128_Causal_FixedLen --bm_name=attention_decode_kv_cache --batch=8 --seq_len_qo=1 --num_heads_q=32 --head_size_qk=128 --head_size_vo=128 --seq_len_kv=1024 --seq_len_kv_cache=1024 --num_heads_kv=8
+PvcFMHADecodeBF16BF16FP32_RCR_h128_NonCausal_FixedLen --bm_name=attention_decode_kv_cache --batch=16 --seq_len_qo=1 --num_heads_q=32 --head_size_qk=128 --head_size_vo=128 --seq_len_kv=1024 --seq_len_kv_cache=1024 --num_heads_kv=8
+PvcFMHADecodeBF16BF16FP32_RCR_h128_Causal_FixedLen --bm_name=attention_decode_kv_cache --batch=16 --seq_len_qo=1 --num_heads_q=32 --head_size_qk=128 --head_size_vo=128 --seq_len_kv=1024 --seq_len_kv_cache=1024 --num_heads_kv=8
+
+PvcFMHADecodeBF16BF16FP32_RCR_h128_NonCausal_VarLen --bm_name=attention_decode_kv_cache --batch=1 --seq_len_qo=1 --num_heads_q=32 --head_size_qk=128 --head_size_vo=128 --seq_len_kv=2048 --seq_len_kv_cache=1024 --num_heads_kv=8
+PvcFMHADecodeBF16BF16FP32_RCR_h128_Causal_VarLen --bm_name=attention_decode_kv_cache --batch=1 --seq_len_qo=1 --num_heads_q=32 --head_size_qk=128 --head_size_vo=128 --seq_len_kv=2048 --seq_len_kv_cache=1024 --num_heads_kv=8
+PvcFMHADecodeBF16BF16FP32_RCR_h128_NonCausal_VarLen --bm_name=attention_decode_kv_cache --batch=8 --seq_len_qo=1 --num_heads_q=32 --head_size_qk=128 --head_size_vo=128 --seq_len_kv=1024 --seq_len_kv_cache=1024 --num_heads_kv=8
+PvcFMHADecodeBF16BF16FP32_RCR_h128_Causal_VarLen --bm_name=attention_decode_kv_cache --batch=8 --seq_len_qo=1 --num_heads_q=32 --head_size_qk=128 --head_size_vo=128 --seq_len_kv=1024 --seq_len_kv_cache=1024 --num_heads_kv=8
+PvcFMHADecodeBF16BF16FP32_RCR_h128_NonCausal_VarLen --bm_name=attention_decode_kv_cache --batch=16 --seq_len_qo=1 --num_heads_q=32 --head_size_qk=128 --head_size_vo=128 --seq_len_kv=1024 --seq_len_kv_cache=1024 --num_heads_kv=8
+PvcFMHADecodeBF16BF16FP32_RCR_h128_Causal_VarLen --bm_name=attention_decode_kv_cache --batch=16 --seq_len_qo=1 --num_heads_q=32 --head_size_qk=128 --head_size_vo=128 --seq_len_kv=1024 --seq_len_kv_cache=1024 --num_heads_kv=8
+
+
+#FP16 benchmarks
+PvcFMHADecodeFP16FP16FP32_RCR_h128_NonCausal_FixedLen --bm_name=attention_decode_kv_cache --batch=1 --seq_len_qo=1 --num_heads_q=32 --head_size_qk=128 --head_size_vo=128 --seq_len_kv=2048 --seq_len_kv_cache=1024 --num_heads_kv=8
+PvcFMHADecodeFP16FP16FP32_RCR_h128_Causal_FixedLen --bm_name=attention_decode_kv_cache --batch=1 --seq_len_qo=1 --num_heads_q=32 --head_size_qk=128 --head_size_vo=128 --seq_len_kv=2048 --seq_len_kv_cache=1024 --num_heads_kv=8
+PvcFMHADecodeFP16FP16FP32_RCR_h128_NonCausal_FixedLen --bm_name=attention_decode_kv_cache --batch=8 --seq_len_qo=1 --num_heads_q=32 --head_size_qk=128 --head_size_vo=128 --seq_len_kv=1024 --seq_len_kv_cache=1024 --num_heads_kv=8
+PvcFMHADecodeFP16FP16FP32_RCR_h128_Causal_FixedLen --bm_name=attention_decode_kv_cache --batch=8 --seq_len_qo=1 --num_heads_q=32 --head_size_qk=128 --head_size_vo=128 --seq_len_kv=1024 --seq_len_kv_cache=1024 --num_heads_kv=8
+PvcFMHADecodeFP16FP16FP32_RCR_h128_NonCausal_FixedLen --bm_name=attention_decode_kv_cache --batch=16 --seq_len_qo=1 --num_heads_q=32 --head_size_qk=128 --head_size_vo=128 --seq_len_kv=1024 --seq_len_kv_cache=1024 --num_heads_kv=8
+PvcFMHADecodeFP16FP16FP32_RCR_h128_Causal_FixedLen --bm_name=attention_decode_kv_cache --batch=16 --seq_len_qo=1 --num_heads_q=32 --head_size_qk=128 --head_size_vo=128 --seq_len_kv=1024 --seq_len_kv_cache=1024 --num_heads_kv=8
+
+PvcFMHADecodeFP16FP16FP32_RCR_h128_NonCausal_VarLen --bm_name=attention_decode_kv_cache --batch=1 --seq_len_qo=1 --num_heads_q=32 --head_size_qk=128 --head_size_vo=128 --seq_len_kv=2048 --seq_len_kv_cache=1024 --num_heads_kv=8
+PvcFMHADecodeFP16FP16FP32_RCR_h128_Causal_VarLen --bm_name=attention_decode_kv_cache --batch=1 --seq_len_qo=1 --num_heads_q=32 --head_size_qk=128 --head_size_vo=128 --seq_len_kv=2048 --seq_len_kv_cache=1024 --num_heads_kv=8
+PvcFMHADecodeFP16FP16FP32_RCR_h128_NonCausal_VarLen --bm_name=attention_decode_kv_cache --batch=8 --seq_len_qo=1 --num_heads_q=32 --head_size_qk=128 --head_size_vo=128 --seq_len_kv=1024 --seq_len_kv_cache=1024 --num_heads_kv=8
+PvcFMHADecodeFP16FP16FP32_RCR_h128_Causal_VarLen --bm_name=attention_decode_kv_cache --batch=8 --seq_len_qo=1 --num_heads_q=32 --head_size_qk=128 --head_size_vo=128 --seq_len_kv=1024 --seq_len_kv_cache=1024 --num_heads_kv=8
+PvcFMHADecodeFP16FP16FP32_RCR_h128_NonCausal_VarLen --bm_name=attention_decode_kv_cache --batch=16 --seq_len_qo=1 --num_heads_q=32 --head_size_qk=128 --head_size_vo=128 --seq_len_kv=1024 --seq_len_kv_cache=1024 --num_heads_kv=8
+PvcFMHADecodeFP16FP16FP32_RCR_h128_Causal_VarLen --bm_name=attention_decode_kv_cache --batch=16 --seq_len_qo=1 --num_heads_q=32 --head_size_qk=128 --head_size_vo=128 --seq_len_kv=1024 --seq_len_kv_cache=1024 --num_heads_kv=8

benchmarks/device/bmg/input_files/input_sglang_flash_attention_decode_nokvcache.in

@@ -0,0 +1,30 @@
+# Flash attention decode (without kv-cache)
+
+PvcFMHADecodeBF16BF16FP32_RCR_h128_NonCausal_FixedLen --bm_name=attention_decode_no_kv_cache --batch=1 --seq_len_qo=1 --num_heads_q=32 --head_size_qk=128 --head_size_vo=128 --seq_len_kv=2048 --seq_len_kv_cache=0 --num_heads_kv=8
+PvcFMHADecodeBF16BF16FP32_RCR_h128_Causal_FixedLen --bm_name=attention_decode_no_kv_cache --batch=1 --seq_len_qo=1 --num_heads_q=32 --head_size_qk=128 --head_size_vo=128 --seq_len_kv=2048 --seq_len_kv_cache=0 --num_heads_kv=8
+PvcFMHADecodeBF16BF16FP32_RCR_h128_NonCausal_FixedLen --bm_name=attention_decode_no_kv_cache --batch=8 --seq_len_qo=1 --num_heads_q=32 --head_size_qk=128 --head_size_vo=128 --seq_len_kv=1024 --seq_len_kv_cache=0 --num_heads_kv=8
+PvcFMHADecodeBF16BF16FP32_RCR_h128_Causal_FixedLen --bm_name=attention_decode_no_kv_cache --batch=8 --seq_len_qo=1 --num_heads_q=32 --head_size_qk=128 --head_size_vo=128 --seq_len_kv=1024 --seq_len_kv_cache=0 --num_heads_kv=8
+PvcFMHADecodeBF16BF16FP32_RCR_h128_NonCausal_FixedLen --bm_name=attention_decode_no_kv_cache --batch=16 --seq_len_qo=1 --num_heads_q=32 --head_size_qk=128 --head_size_vo=128 --seq_len_kv=1024 --seq_len_kv_cache=0 --num_heads_kv=8
+PvcFMHADecodeBF16BF16FP32_RCR_h128_Causal_FixedLen --bm_name=attention_decode_no_kv_cache --batch=16 --seq_len_qo=1 --num_heads_q=32 --head_size_qk=128 --head_size_vo=128 --seq_len_kv=1024 --seq_len_kv_cache=0 --num_heads_kv=8
+
+PvcFMHADecodeBF16BF16FP32_RCR_h128_NonCausal_VarLen --bm_name=attention_decode_no_kv_cache --batch=1 --seq_len_qo=1 --num_heads_q=32 --head_size_qk=128 --head_size_vo=128 --seq_len_kv=2048 --seq_len_kv_cache=0 --num_heads_kv=8
+PvcFMHADecodeBF16BF16FP32_RCR_h128_Causal_VarLen --bm_name=attention_decode_no_kv_cache --batch=1 --seq_len_qo=1 --num_heads_q=32 --head_size_qk=128 --head_size_vo=128 --seq_len_kv=2048 --seq_len_kv_cache=0 --num_heads_kv=8
+PvcFMHADecodeBF16BF16FP32_RCR_h128_NonCausal_VarLen --bm_name=attention_decode_no_kv_cache --batch=8 --seq_len_qo=1 --num_heads_q=32 --head_size_qk=128 --head_size_vo=128 --seq_len_kv=1024 --seq_len_kv_cache=0 --num_heads_kv=8
+PvcFMHADecodeBF16BF16FP32_RCR_h128_Causal_VarLen --bm_name=attention_decode_no_kv_cache --batch=8 --seq_len_qo=1 --num_heads_q=32 --head_size_qk=128 --head_size_vo=128 --seq_len_kv=1024 --seq_len_kv_cache=0 --num_heads_kv=8
+PvcFMHADecodeBF16BF16FP32_RCR_h128_NonCausal_VarLen --bm_name=attention_decode_no_kv_cache --batch=16 --seq_len_qo=1 --num_heads_q=32 --head_size_qk=128 --head_size_vo=128 --seq_len_kv=1024 --seq_len_kv_cache=0 --num_heads_kv=8
+PvcFMHADecodeBF16BF16FP32_RCR_h128_Causal_VarLen --bm_name=attention_decode_no_kv_cache --batch=16 --seq_len_qo=1 --num_heads_q=32 --head_size_qk=128 --head_size_vo=128 --seq_len_kv=1024 --seq_len_kv_cache=0 --num_heads_kv=8
+
+#FP16 benchmarks
+PvcFMHADecodeFP16FP16FP32_RCR_h128_NonCausal_FixedLen --bm_name=attention_decode_no_kv_cache --batch=1 --seq_len_qo=1 --num_heads_q=32 --head_size_qk=128 --head_size_vo=128 --seq_len_kv=2048 --seq_len_kv_cache=0 --num_heads_kv=8
+PvcFMHADecodeFP16FP16FP32_RCR_h128_Causal_FixedLen --bm_name=attention_decode_no_kv_cache --batch=1 --seq_len_qo=1 --num_heads_q=32 --head_size_qk=128 --head_size_vo=128 --seq_len_kv=2048 --seq_len_kv_cache=0 --num_heads_kv=8
+PvcFMHADecodeFP16FP16FP32_RCR_h128_NonCausal_FixedLen --bm_name=attention_decode_no_kv_cache --batch=8 --seq_len_qo=1 --num_heads_q=32 --head_size_qk=128 --head_size_vo=128 --seq_len_kv=1024 --seq_len_kv_cache=0 --num_heads_kv=8
+PvcFMHADecodeFP16FP16FP32_RCR_h128_Causal_FixedLen --bm_name=attention_decode_no_kv_cache --batch=8 --seq_len_qo=1 --num_heads_q=32 --head_size_qk=128 --head_size_vo=128 --seq_len_kv=1024 --seq_len_kv_cache=0 --num_heads_kv=8
+PvcFMHADecodeFP16FP16FP32_RCR_h128_NonCausal_FixedLen --bm_name=attention_decode_no_kv_cache --batch=16 --seq_len_qo=1 --num_heads_q=32 --head_size_qk=128 --head_size_vo=128 --seq_len_kv=1024 --seq_len_kv_cache=0 --num_heads_kv=8
+PvcFMHADecodeFP16FP16FP32_RCR_h128_Causal_FixedLen --bm_name=attention_decode_no_kv_cache --batch=16 --seq_len_qo=1 --num_heads_q=32 --head_size_qk=128 --head_size_vo=128 --seq_len_kv=1024 --seq_len_kv_cache=0 --num_heads_kv=8
+
+PvcFMHADecodeFP16FP16FP32_RCR_h128_NonCausal_VarLen --bm_name=attention_decode_no_kv_cache --batch=1 --seq_len_qo=1 --num_heads_q=32 --head_size_qk=128 --head_size_vo=128 --seq_len_kv=2048 --seq_len_kv_cache=0 --num_heads_kv=8
+PvcFMHADecodeFP16FP16FP32_RCR_h128_Causal_VarLen --bm_name=attention_decode_no_kv_cache --batch=1 --seq_len_qo=1 --num_heads_q=32 --head_size_qk=128 --head_size_vo=128 --seq_len_kv=2048 --seq_len_kv_cache=0 --num_heads_kv=8
+PvcFMHADecodeFP16FP16FP32_RCR_h128_NonCausal_VarLen --bm_name=attention_decode_no_kv_cache --batch=8 --seq_len_qo=1 --num_heads_q=32 --head_size_qk=128 --head_size_vo=128 --seq_len_kv=1024 --seq_len_kv_cache=0 --num_heads_kv=8
+PvcFMHADecodeFP16FP16FP32_RCR_h128_Causal_VarLen --bm_name=attention_decode_no_kv_cache --batch=8 --seq_len_qo=1 --num_heads_q=32 --head_size_qk=128 --head_size_vo=128 --seq_len_kv=1024 --seq_len_kv_cache=0 --num_heads_kv=8
+PvcFMHADecodeFP16FP16FP32_RCR_h128_NonCausal_VarLen --bm_name=attention_decode_no_kv_cache --batch=16 --seq_len_qo=1 --num_heads_q=32 --head_size_qk=128 --head_size_vo=128 --seq_len_kv=1024 --seq_len_kv_cache=0 --num_heads_kv=8
+PvcFMHADecodeFP16FP16FP32_RCR_h128_Causal_VarLen --bm_name=attention_decode_no_kv_cache --batch=16 --seq_len_qo=1 --num_heads_q=32 --head_size_qk=128 --head_size_vo=128 --seq_len_kv=1024 --seq_len_kv_cache=0 --num_heads_kv=8

benchmarks/flash_attention/CMakeLists.txt

@@ -28,3 +28,4 @@
 
 cutlass_benchmark_add_suite(cutlass_benchmarks_flash_attention)
 add_subdirectory(flash_attention_prefill)
+add_subdirectory(flash_attention_decode)

benchmarks/flash_attention/flash_attention_decode/CMakeLists.txt

@@ -0,0 +1,45 @@
+# Copyright (c) 2024 - 2025 Codeplay Software Ltd. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+#
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+#
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+set(CUTLASS_APPLICATIONS_DIR ${CMAKE_SOURCE_DIR}/applications)
+
+# Pass these configuration files for the CI
+set(CONFIG_FILE_KV_CACHE --config_file=${CMAKE_SOURCE_DIR}/benchmarks/device/bmg/input_files/input_sglang_flash_attention_decode_kvcache.in)
+set(CONFIG_FILE_NO_KV_CACHE --config_file=${CMAKE_SOURCE_DIR}/benchmarks/device/bmg/input_files/input_sglang_flash_attention_decode_nokvcache.in)
+
+cutlass_benchmark_add_suite(cutlass_benchmarks_flash_attention_decode
+                            SUPERSUITE cutlass_benchmarks_flash_attention)
+
+cutlass_benchmark_add_executable(
+    cutlass_benchmarks_flash_attention_decode_xe
+    main.cpp
+    TEST_COMMAND_OPTIONS
+    CONFIG_FILE_NO_KV_CACHE
+    CONFIG_FILE_KV_CACHE
+    SUITE cutlass_benchmarks_flash_attention_decode
+)