Add base framework for FlashAttention unit tests (#295)

This PR adds the foundational framework for unit testing FlashAttention.
While it doesn't cover all test scenarios it provides a starting point
for validating core functionality and ensuring correctness.

Author: aacostadiaz

applications/flash_attention_v2/collective/fmha_fusion.hpp

@@ -0,0 +1,95 @@
+/***************************************************************************************************
+* Copyright (c) 2025 - 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.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include "cutlass/cutlass.h"
+#include <cute/tensor.hpp>
+
+namespace cutlass::fmha::collective {
+
+using namespace cute;
+
+struct VariableLength {
+  int max_length;
+  int* cumulative_length = nullptr;
+
+  CUTE_HOST_DEVICE operator int() const {
+    return max_length;
+  }
+};
+
+template<class T> struct is_variable_length : std::false_type {};
+template<> struct is_variable_length<VariableLength> : std::true_type {};
+template<class T> constexpr bool is_variable_length_v = is_variable_length<T>::value;
+
+template<class Shape, class Idx>
+CUTE_HOST_DEVICE
+constexpr auto
+apply_variable_length(Shape const& shape, Idx const& idx) {
+  return transform_leaf(shape, [&](auto const& s) {
+    if constexpr (is_variable_length_v<remove_cvref_t<decltype(s)>>) {
+      return s.cumulative_length[idx+1] - s.cumulative_length[idx];
+    }
+    else {
+      return s;
+    }
+  });
+}
+
+template<class Shape, class Coord, class Idx>
+CUTE_HOST_DEVICE
+constexpr auto
+apply_variable_length(Shape const& shape, Coord const& coord, Idx const& idx) {
+  auto new_shape = apply_variable_length(shape, idx);
+  auto new_coord = transform_leaf(shape, coord, [&](auto const& s, auto const& c) {
+    if constexpr (is_variable_length_v<remove_cvref_t<decltype(s)>>) {
+      return cute::make_tuple(c, s.cumulative_length[idx]);
+    }
+    else {
+      return c;
+    }
+  });
+  return cute::make_tuple(new_shape, new_coord);
+}
+
+}  // namespace cutlass::fmha::collective
+
+namespace cute {
+
+template<>
+struct is_integral<cutlass::fmha::collective::VariableLength> : true_type {};
+
+CUTE_HOST_DEVICE
+void print(cutlass::fmha::collective::VariableLength a) {
+  printf("Varlen<%d, %p>", a.max_length, a.cumulative_length);
+}
+
+}

applications/flash_attention_v2/collective/xe_flash_attn_mma.hpp

@@ -36,7 +36,9 @@
 #include "cute/algorithm/functional.hpp"
 #include "cute/atom/mma_atom.hpp"
 #include "cute/algorithm/gemm.hpp"
-#include "cute/tensor_predicate.hpp"
+#include "cutlass/util/packed_stride.hpp"
+
+#include "fmha_fusion.hpp"
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 

applications/flash_attention_v2/kernel/xe_flash_attn_gemm.hpp

@@ -36,6 +36,7 @@
 #include "cutlass/kernel_hardware_info.hpp"
 
 #include "flash_attention_v2/collective/xe_flash_attn_mma.hpp"
+#include "flash_attention_v2/kernel/tile_scheduler.hpp"
 
 namespace cutlass::flash_attention::kernel {
 

benchmarks/pvc/flash_attention_v2/benchmark_runner.hpp

@@ -32,7 +32,6 @@
 
 #include "cutlass/epilogue/collective/default_epilogue.hpp"
 #include "cutlass/epilogue/fusion/xe_callbacks.hpp"
-#include "77_blackwell_fmha/collective/fmha_fusion.hpp"
 #include "flash_attention_v2/kernel/tile_scheduler.hpp"
 #include "cutlass/gemm/device/gemm_universal_adapter.h"
 #include "cutlass/util/packed_stride.hpp"

examples/sycl/06_pvc_flash_attention/pvc_flash_attn_runner.hpp

@@ -32,7 +32,7 @@
 
 #include "cutlass/epilogue/collective/default_epilogue.hpp"
 #include "cutlass/gemm/device/gemm_universal_adapter.h"
-#include "77_blackwell_fmha/collective/fmha_fusion.hpp"
+#include "flash_attention_v2/collective/fmha_fusion.hpp"
 #include "flash_attention_v2/kernel/tile_scheduler.hpp"
 #include "cutlass/util/packed_stride.hpp"
 #include "flash_attention_v2/kernel/xe_flash_attn_gemm.hpp"

test/unit/CMakeLists.txt

@@ -85,6 +85,7 @@ function(cutlass_test_unit_add_executable NAME)
     PRIVATE
     ${CUTLASS_UNIT_TEST_COMMON_DIR}
     ${__EXTRA_INCLUDE_DIRS}
+    ${CUTLASS_APPLICATIONS_DIR}
   )
   if (__WITHOUT_CUDA)
     # Avoid CUDA dependencies for host-only unit tests that provide the
@@ -108,6 +109,7 @@ function(cutlass_test_unit_add_executable NAME)
   endif()
 
   if(CUTLASS_ENABLE_SYCL)
+    add_onemkl_to_target(TARGET ${NAME})
     add_sycl_to_target(TARGET ${NAME})
   endif()
 
@@ -188,6 +190,7 @@ if (CUTLASS_ENABLE_SYCL)
   set(SUBDIRS
     cute
     gemm
+    flash_attention
   )
 else()
   set(SUBDIRS

test/unit/flash_attention/CMakeLists.txt

@@ -0,0 +1,56 @@
+# Copyright (c) 2025 - 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.
+
+if(SYCL_INTEL_TARGET)
+
+set(CUTLASS_APPLICATIONS_DIR ${CMAKE_SOURCE_DIR}/applications)
+
+cutlass_test_unit_add_executable(
+  cutlass_test_unit_flash_attention_xe
+  xe_flash_attention_bf16.cpp
+  xe_flash_attention_fp16.cpp
+)
+
+add_custom_target(
+  cutlass_test_unit_flash_attention
+  DEPENDS
+  cutlass_test_unit_flash_attention_xe
+)
+
+add_custom_target(
+  test_unit_flash_attention
+  DEPENDS
+  test_unit_flash_attention_xe
+)
+
+else()
+  # Dummy targets if not building for Intel
+  add_custom_target(cutlass_test_unit_flash_attention)
+  add_custom_target(test_unit_flash_attention)
+
+endif()