Add varlen attention support (no vmap or sharding support yet).

Author: nshepperd

.clang-format

@@ -0,0 +1,4 @@
+---
+Language: Cpp
+ColumnLimit: 100
+BasedOnStyle: Google

CMakeLists.txt

@@ -27,7 +27,7 @@ string(REGEX REPLACE "--generate-code=arch=compute_[0-9]+,code=\\[?compute_[0-9]
 string(REGEX REPLACE "-gencode arch=compute_[0-9]+,code=sm_[0-9]+" "" 
        CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS}")
 
-message(WARNING "CMAKE_CUDA_FLAGS: ${CMAKE_CUDA_FLAGS}")
+message(STATUS "CMAKE_CUDA_FLAGS: ${CMAKE_CUDA_FLAGS}")
 
 # Set up ccache
 find_program(CCACHE_PROGRAM ccache)
@@ -49,9 +49,11 @@ find_package(CUDAToolkit REQUIRED)
 
 
 # CUDA flags
+set(CMAKE_CXX_STANDARD 20)
+set(CMAKE_CXX_STANDARD_REQUIRED ON)
+set(CMAKE_CXX_EXTENSIONS OFF)
 set(CUDA_FLAGS
     -O3
-    -std=c++20
     --use_fast_math
     --expt-relaxed-constexpr
     --expt-extended-lambda

csrc/flash_attn/check.h

@@ -14,7 +14,6 @@ class CheckHelper {
   explicit CheckHelper(std::string expr) : expr_(expr) {}
 
   template <typename T> inline CheckHelper &operator<<(const T &value) {
-    fprintf(stderr, "debug: adding value %s\n", value);
     stream_ << value;
     return *this;
   }
@@ -29,7 +28,6 @@ class CheckHelper {
     full_message << "Check failed: " << expr_;
     std::string additional = stream_.str();
     if (!additional.empty()) {
-      fprintf(stderr, "debug: %s\n", additional.c_str());
       full_message << "; " << additional;
     }
     return ffi::Error(errc_, full_message.str());
@@ -63,3 +61,10 @@ class CheckHelper {
 #define FFI_RET_CHECK(expr)                                                    \
   if (auto _error = (expr); !_error.success())                                 \
   return _error
+
+#define FFI_CHECK_ALLOC(dest, expr)                                         \
+  void* dest = nullptr;                                  \ 
+  if (auto _opt = (expr); _opt.has_value())                                    \
+    dest = _opt.value();                                                       \
+  else                                                                         \
+    return CheckHelper(std::string(#expr))

csrc/flash_attn/flash_api.cpp

@@ -338,11 +338,61 @@ XLA_FFI_DEFINE_HANDLER(
 		.Attr<int64_t>("window_size_right")
 );
 
+XLA_FFI_DEFINE_HANDLER(
+	mha_varlen_fwd, mha_varlen_fwd_impl,
+	ffi::Ffi::Bind()
+		.Ctx<ffi::PlatformStream<cudaStream_t>>()
+		.Ctx<ffi::ScratchAllocator>()
+		.Arg<ffi::AnyBuffer>() // q
+		.Arg<ffi::AnyBuffer>() // k
+		.Arg<ffi::AnyBuffer>() // v
+		.Arg<ffi::Buffer<ffi::S32>>() // cu_seqlens_q
+		.Arg<ffi::Buffer<ffi::S32>>() // cu_seqlens_k
+		.Arg<ffi::Buffer<ffi::S32>>() // seqused_k
+		.Ret<ffi::AnyBuffer>() // o
+		.Ret<ffi::Buffer<ffi::F32>>() // lse
+		.Attr<int>("max_seqlen_q")
+		.Attr<int>("max_seqlen_k")
+		.Attr<bool>("has_seqused_k")
+		.Attr<double>("softmax_scale")
+		.Attr<bool>("zero_tensors")
+		.Attr<bool>("is_causal")
+		.Attr<int64_t>("window_size_left")
+		.Attr<int64_t>("window_size_right")
+);
+
+XLA_FFI_DEFINE_HANDLER(
+	mha_varlen_bwd, mha_varlen_bwd_impl,
+	ffi::Ffi::Bind()
+		.Ctx<ffi::PlatformStream<cudaStream_t>>()
+		.Ctx<ffi::ScratchAllocator>()
+		.Arg<ffi::AnyBuffer>() // dout
+		.Arg<ffi::AnyBuffer>() // q
+		.Arg<ffi::AnyBuffer>() // k
+		.Arg<ffi::AnyBuffer>() // v
+		.Arg<ffi::AnyBuffer>() // o
+		.Arg<ffi::Buffer<ffi::F32>>() // lse
+		.Arg<ffi::Buffer<ffi::S32>>() // cu_seqlens_q
+		.Arg<ffi::Buffer<ffi::S32>>() // cu_seqlens_k
+		.Ret<ffi::AnyBuffer>() // dq
+		.Ret<ffi::AnyBuffer>() // dk
+		.Ret<ffi::AnyBuffer>() // dv
+		.Attr<int64_t>("max_seqlen_q")
+		.Attr<int64_t>("max_seqlen_k")
+		.Attr<float>("softmax_scale")
+		.Attr<bool>("zero_tensors")
+		.Attr<bool>("is_causal")
+		.Attr<int64_t>("window_size_left")
+		.Attr<int64_t>("window_size_right")
+		.Attr<bool>("deterministic")
+);
 
 pybind11::dict FFIRegistrations() {
   pybind11::dict dict;
   dict["flash_mha_fwd"] = EncapsulateFfiCall(mha_fwd);
   dict["flash_mha_bwd"] = EncapsulateFfiCall(mha_bwd);
+  dict["flash_mha_varlen_fwd"] = EncapsulateFfiCall(mha_varlen_fwd);
+  dict["flash_mha_varlen_bwd"] = EncapsulateFfiCall(mha_varlen_bwd);
   return dict;
 }