feat(linked): add FlashAttention FlashAttnVarlenFunc provider

src/base/flash_attn_varlen_func.h

@@ -52,6 +52,12 @@ class FlashAttnVarlenFunc : public Operator<FlashAttnVarlenFunc> {
         v_shape_{v.shape()},
         cu_seqlens_q_shape_{cu_seqlens_q.shape()},
         cu_seqlens_k_shape_{cu_seqlens_k.shape()},
+        alibi_slopes_shape_{alibi_slopes.has_value()
+                                ? Tensor::Shape{alibi_slopes->shape()}
+                                : Tensor::Shape{}},
+        block_table_shape_{block_table.has_value()
+                               ? Tensor::Shape{block_table->shape()}
+                               : Tensor::Shape{}},
         out_shape_{out.shape()},
         softmax_lse_shape_{softmax_lse.has_value()
                                ? Tensor::Shape{softmax_lse->shape()}
@@ -63,6 +69,12 @@ class FlashAttnVarlenFunc : public Operator<FlashAttnVarlenFunc> {
         v_strides_{v.strides()},
         cu_seqlens_q_strides_{cu_seqlens_q.strides()},
         cu_seqlens_k_strides_{cu_seqlens_k.strides()},
+        alibi_slopes_strides_{alibi_slopes.has_value()
+                                  ? Tensor::Strides{alibi_slopes->strides()}
+                                  : Tensor::Strides{}},
+        block_table_strides_{block_table.has_value()
+                                 ? Tensor::Strides{block_table->strides()}
+                                 : Tensor::Strides{}},
         out_strides_{out.strides()},
         softmax_lse_strides_{softmax_lse.has_value()
                                  ? Tensor::Strides{softmax_lse->strides()}
@@ -75,18 +87,25 @@ class FlashAttnVarlenFunc : public Operator<FlashAttnVarlenFunc> {
         v_dtype_{v.dtype()},
         cu_seqlens_q_dtype_{cu_seqlens_q.dtype()},
         cu_seqlens_k_dtype_{cu_seqlens_k.dtype()},
+        alibi_slopes_dtype_{alibi_slopes.has_value() ? alibi_slopes->dtype()
+                                                     : DataType::kFloat32},
+        block_table_dtype_{block_table.has_value() ? block_table->dtype()
+                                                   : DataType::kInt32},
         out_dtype_{out.dtype()},
         softmax_lse_dtype_{softmax_lse.has_value() ? softmax_lse->dtype()
                                                    : DataType::kFloat32},
         s_dmask_dtype_{s_dmask.has_value() ? s_dmask->dtype() : q.dtype()},
         has_auxiliary_outputs_{softmax_lse.has_value() && s_dmask.has_value()},
         device_index_{q.device().index()} {
-    assert(q.ndim() == 3 && k.ndim() == 3 && v.ndim() == 3 &&
-           "`FlashAttnVarlenFunc` requires packed 3D Q, K, and V tensors");
+    assert(q.ndim() == 3 &&
+           ((!block_table.has_value() && k.ndim() == 3 && v.ndim() == 3) ||
+            (block_table.has_value() && k.ndim() == 4 && v.ndim() == 4)) &&
+           "`FlashAttnVarlenFunc` requires packed 3D Q and either packed 3D "
+           "or paged 4D K and V tensors");
     assert(k.shape() == v.shape() &&
            "`FlashAttnVarlenFunc` requires K and V to have the same shape");
-    assert(q.size(1) > 0 && k.size(1) > 0 && q.size(2) == k.size(2) &&
-           q.size(1) % k.size(1) == 0 &&
+    assert(q.size(1) > 0 && k.size(-2) > 0 && q.size(2) == k.size(-1) &&
+           q.size(1) % k.size(-2) == 0 &&
            "`FlashAttnVarlenFunc` requires compatible Q and KV heads");
     assert(q.size(2) > 0 && q.size(2) <= 256 && q.size(2) % 8 == 0 &&
            "`FlashAttnVarlenFunc` requires a head dimension divisible by 8 "
@@ -140,10 +159,25 @@ class FlashAttnVarlenFunc : public Operator<FlashAttnVarlenFunc> {
            "`FlashAttnVarlenFunc` does not yet support softcap");
     assert(!deterministic &&
            "`FlashAttnVarlenFunc` does not yet support deterministic mode");
-    assert(!block_table.has_value() &&
-           "`FlashAttnVarlenFunc` does not yet support paged KV cache");
-    assert(!alibi_slopes.has_value() &&
-           "`FlashAttnVarlenFunc` does not yet support ALiBi slopes");
+    if (block_table.has_value()) {
+      assert(block_table->ndim() == 2 &&
+             block_table->size(0) + 1 == cu_seqlens_q.size(0) &&
+             block_table_dtype_ == DataType::kInt32 &&
+             block_table->IsContiguous() && k.size(1) % 256 == 0 &&
+             "`FlashAttnVarlenFunc` requires a contiguous int32 block table "
+             "and page size divisible by 256");
+    }
+    if (alibi_slopes.has_value()) {
+      assert(
+          (alibi_slopes->ndim() == 1 || alibi_slopes->ndim() == 2) &&
+          alibi_slopes_dtype_ == DataType::kFloat32 &&
+          alibi_slopes->IsContiguous() &&
+          ((alibi_slopes->ndim() == 1 && alibi_slopes->size(0) == q.size(1)) ||
+           (alibi_slopes->ndim() == 2 &&
+            alibi_slopes->size(0) + 1 == cu_seqlens_q.size(0) &&
+            alibi_slopes->size(1) == q.size(1))) &&
+          "`FlashAttnVarlenFunc` received incompatible ALiBi slopes");
+    }
 
     const auto same_device_as_q = [&](const Tensor tensor) {
       return tensor.device().type() == q.device().type() &&
@@ -152,6 +186,8 @@ class FlashAttnVarlenFunc : public Operator<FlashAttnVarlenFunc> {
     assert(same_device_as_q(k) && same_device_as_q(v) &&
            same_device_as_q(cu_seqlens_q) && same_device_as_q(cu_seqlens_k) &&
            same_device_as_q(out) &&
+           (!alibi_slopes.has_value() || same_device_as_q(*alibi_slopes)) &&
+           (!block_table.has_value() || same_device_as_q(*block_table)) &&
            (!softmax_lse.has_value() || same_device_as_q(*softmax_lse)) &&
            (!s_dmask.has_value() || same_device_as_q(*s_dmask)) &&
            "`FlashAttnVarlenFunc` tensors must be on the same device");
@@ -191,6 +227,10 @@ class FlashAttnVarlenFunc : public Operator<FlashAttnVarlenFunc> {
 
   Tensor::Shape cu_seqlens_k_shape_;
 
+  Tensor::Shape alibi_slopes_shape_;
+
+  Tensor::Shape block_table_shape_;
+
   Tensor::Shape out_shape_;
 
   Tensor::Shape softmax_lse_shape_;
@@ -207,6 +247,10 @@ class FlashAttnVarlenFunc : public Operator<FlashAttnVarlenFunc> {
 
   Tensor::Strides cu_seqlens_k_strides_;
 
+  Tensor::Strides alibi_slopes_strides_;
+
+  Tensor::Strides block_table_strides_;
+
   Tensor::Strides out_strides_;
 
   Tensor::Strides softmax_lse_strides_;
@@ -223,6 +267,10 @@ class FlashAttnVarlenFunc : public Operator<FlashAttnVarlenFunc> {
 
   DataType cu_seqlens_k_dtype_;
 
+  DataType alibi_slopes_dtype_;
+
+  DataType block_table_dtype_;
+
   DataType out_dtype_;
 
   DataType softmax_lse_dtype_;

src/linked/torch/nvidia/ops/flash_attn_varlen_func/flash_attn.cc

@@ -0,0 +1,44 @@
+#include "linked/torch/nvidia/ops/flash_attn_varlen_func/flash_attn.h"
+
+namespace flash {
+
+std::vector<at::Tensor> mha_varlen_fwd(
+    at::Tensor& q, const at::Tensor& k, const at::Tensor& v,
+    std::optional<at::Tensor>& out, const at::Tensor& cu_seqlens_q,
+    const at::Tensor& cu_seqlens_k, std::optional<at::Tensor>& seqused_k,
+    std::optional<const at::Tensor>& leftpad_k,
+    std::optional<at::Tensor>& block_table,
+    std::optional<at::Tensor>& alibi_slopes, int max_seqlen_q, int max_seqlen_k,
+    float dropout_p, float softmax_scale, bool zero_tensors, bool causal,
+    int window_size_left, int window_size_right, float softcap,
+    bool return_softmax, std::optional<at::Generator> generator);
+
+}  // namespace flash
+
+namespace infini::ops::linked::torch::nvidia {
+
+std::vector<at::Tensor> FlashAttnVarlen::Call(
+    at::Tensor& q, const at::Tensor& k, const at::Tensor& v,
+    std::optional<at::Tensor>& out, const at::Tensor& cu_seqlens_q,
+    const at::Tensor& cu_seqlens_k, std::optional<at::Tensor>& seqused_k,
+    std::optional<const at::Tensor>& leftpad_k,
+    std::optional<at::Tensor>& block_table,
+    std::optional<at::Tensor>& alibi_slopes, int max_seqlen_q, int max_seqlen_k,
+    float dropout_p, float softmax_scale, bool zero_tensors, bool causal,
+    int window_size_left, int window_size_right, float softcap,
+    bool return_softmax, std::optional<at::Generator> generator) {
+  return flash::mha_varlen_fwd(
+      q, k, v, out, cu_seqlens_q, cu_seqlens_k, seqused_k, leftpad_k,
+      block_table, alibi_slopes, max_seqlen_q, max_seqlen_k, dropout_p,
+      softmax_scale, zero_tensors, causal, window_size_left, window_size_right,
+      softcap, return_softmax, generator);
+}
+
+}  // namespace infini::ops::linked::torch::nvidia
+
+namespace infini::ops::linked::torch {
+
+template class TorchFlashAttnVarlenFunc<
+    ::infini::ops::linked::torch::nvidia::FlashAttnVarlen>;
+
+}  // namespace infini::ops::linked::torch

src/linked/torch/nvidia/ops/flash_attn_varlen_func/flash_attn.h

@@ -0,0 +1,49 @@
+#ifndef INFINI_OPS_LINKED_TORCH_NVIDIA_OPS_FLASH_ATTN_VARLEN_FUNC_FLASH_ATTN_H_
+#define INFINI_OPS_LINKED_TORCH_NVIDIA_OPS_FLASH_ATTN_VARLEN_FUNC_FLASH_ATTN_H_
+
+#include <ATen/core/Generator.h>
+
+#include "linked/torch/nvidia/c10.h"
+#include "linked/torch/ops/flash_attn_varlen_func.h"
+
+namespace infini::ops::linked::torch::nvidia {
+
+struct FlashAttnVarlen : C10<Device::Type::kNvidia> {
+  static std::vector<at::Tensor> Call(
+      at::Tensor& q, const at::Tensor& k, const at::Tensor& v,
+      std::optional<at::Tensor>& out, const at::Tensor& cu_seqlens_q,
+      const at::Tensor& cu_seqlens_k, std::optional<at::Tensor>& seqused_k,
+      std::optional<const at::Tensor>& leftpad_k,
+      std::optional<at::Tensor>& block_table,
+      std::optional<at::Tensor>& alibi_slopes, int max_seqlen_q,
+      int max_seqlen_k, float dropout_p, float softmax_scale, bool zero_tensors,
+      bool causal, int window_size_left, int window_size_right, float softcap,
+      bool return_softmax, std::optional<at::Generator> generator);
+};
+
+}  // namespace infini::ops::linked::torch::nvidia
+
+namespace infini::ops::linked::torch {
+
+extern template class TorchFlashAttnVarlenFunc<
+    ::infini::ops::linked::torch::nvidia::FlashAttnVarlen>;
+
+}  // namespace infini::ops::linked::torch
+
+namespace infini::ops {
+
+template <>
+class Operator<FlashAttnVarlenFunc, Device::Type::kNvidia, 16>
+    : public linked::torch::TorchFlashAttnVarlenFunc<
+          linked::torch::nvidia::FlashAttnVarlen> {
+ public:
+  using linked::torch::TorchFlashAttnVarlenFunc<
+      linked::torch::nvidia::FlashAttnVarlen>::TorchFlashAttnVarlenFunc;
+
+  using linked::torch::TorchFlashAttnVarlenFunc<
+      linked::torch::nvidia::FlashAttnVarlen>::operator();
+};
+
+}  // namespace infini::ops
+
+#endif  // INFINI_OPS_LINKED_TORCH_NVIDIA_OPS_FLASH_ATTN_VARLEN_FUNC_FLASH_ATTN_H_

src/linked/torch/nvidia/ops/flash_attn_varlen_func/flash_attn.yaml

@@ -0,0 +1,4 @@
+library: flash_attn
+required_symbols:
+  - >-
+    flash::mha_varlen_fwd(at::Tensor&, at::Tensor const&, at::Tensor const&, std::optional<at::Tensor>&, at::Tensor const&, at::Tensor const&, std::optional<at::Tensor>&, std::optional<at::Tensor const>&, std::optional<at::Tensor>&, std::optional<at::Tensor>&, int, int, float, float, bool, bool, int, int, float, bool, std::optional<at::Generator>)

src/linked/torch/ops/flash_attn_varlen_func.h

@@ -0,0 +1,110 @@
+#ifndef INFINI_OPS_LINKED_TORCH_OPS_FLASH_ATTN_VARLEN_FUNC_H_
+#define INFINI_OPS_LINKED_TORCH_OPS_FLASH_ATTN_VARLEN_FUNC_H_
+
+#include <ATen/core/Generator.h>
+
+#include <cassert>
+#include <cmath>
+#include <optional>
+#include <vector>
+
+#include "base/flash_attn_varlen_func.h"
+#include "torch/tensor_.h"
+
+namespace infini::ops::linked::torch {
+
+template <typename Backend>
+class TorchFlashAttnVarlenFunc : public ::infini::ops::FlashAttnVarlenFunc {
+ public:
+  using ::infini::ops::FlashAttnVarlenFunc::FlashAttnVarlenFunc;
+
+  using ::infini::ops::FlashAttnVarlenFunc::operator();
+
+  void operator()(const Tensor q, const Tensor k, const Tensor v,
+                  const Tensor cu_seqlens_q, const Tensor cu_seqlens_k,
+                  const std::optional<Tensor> alibi_slopes,
+                  const std::optional<Tensor> block_table,
+                  const int64_t max_seqlen_q, const int64_t max_seqlen_k,
+                  const double dropout_p,
+                  const std::optional<double> softmax_scale, const bool causal,
+                  const std::vector<int64_t> window_size, const double softcap,
+                  const bool deterministic, const bool return_attn_probs,
+                  Tensor out, std::optional<Tensor> softmax_lse,
+                  std::optional<Tensor> s_dmask) const override {
+    const typename Backend::StreamGuard stream_guard{
+        Backend::GetStreamFromExternal(stream_, device_index_)};
+
+    auto at_q = ToAtenTensor<Backend::kDeviceType>(const_cast<void*>(q.data()),
+                                                   q_shape_, q_strides_,
+                                                   q_dtype_, device_index_);
+    auto at_k = ToAtenTensor<Backend::kDeviceType>(const_cast<void*>(k.data()),
+                                                   k_shape_, k_strides_,
+                                                   k_dtype_, device_index_);
+    auto at_v = ToAtenTensor<Backend::kDeviceType>(const_cast<void*>(v.data()),
+                                                   v_shape_, v_strides_,
+                                                   v_dtype_, device_index_);
+    auto at_cu_seqlens_q = ToAtenTensor<Backend::kDeviceType>(
+        const_cast<void*>(cu_seqlens_q.data()), cu_seqlens_q_shape_,
+        cu_seqlens_q_strides_, cu_seqlens_q_dtype_, device_index_);
+    auto at_cu_seqlens_k = ToAtenTensor<Backend::kDeviceType>(
+        const_cast<void*>(cu_seqlens_k.data()), cu_seqlens_k_shape_,
+        cu_seqlens_k_strides_, cu_seqlens_k_dtype_, device_index_);
+    auto at_out = ToAtenTensor<Backend::kDeviceType>(
+        out.data(), out_shape_, out_strides_, out_dtype_, device_index_);
+
+    std::optional<at::Tensor> at_alibi_slopes;
+    std::optional<at::Tensor> at_block_table;
+    std::optional<at::Tensor> at_softmax_lse;
+    std::optional<at::Tensor> at_s_dmask;
+
+    if (alibi_slopes.has_value()) {
+      at_alibi_slopes.emplace(ToAtenTensor<Backend::kDeviceType>(
+          const_cast<void*>(alibi_slopes->data()), alibi_slopes_shape_,
+          alibi_slopes_strides_, alibi_slopes_dtype_, device_index_));
+    }
+    if (block_table.has_value()) {
+      at_block_table.emplace(ToAtenTensor<Backend::kDeviceType>(
+          const_cast<void*>(block_table->data()), block_table_shape_,
+          block_table_strides_, block_table_dtype_, device_index_));
+    }
+    if (softmax_lse.has_value()) {
+      at_softmax_lse.emplace(ToAtenTensor<Backend::kDeviceType>(
+          softmax_lse->data(), softmax_lse_shape_, softmax_lse_strides_,
+          softmax_lse_dtype_, device_index_));
+      at_s_dmask.emplace(ToAtenTensor<Backend::kDeviceType>(
+          s_dmask->data(), s_dmask_shape_, s_dmask_strides_, s_dmask_dtype_,
+          device_index_));
+    }
+
+    std::optional<at::Tensor> at_out_optional;
+    std::optional<at::Tensor> at_seqused_k;
+    std::optional<const at::Tensor> at_leftpad_k;
+    std::optional<at::Generator> generator;
+    auto result = Backend::Call(
+        at_q, at_k, at_v, at_out_optional, at_cu_seqlens_q, at_cu_seqlens_k,
+        at_seqused_k, at_leftpad_k, at_block_table, at_alibi_slopes,
+        static_cast<int>(max_seqlen_q), static_cast<int>(max_seqlen_k),
+        static_cast<float>(dropout_p),
+        static_cast<float>(softmax_scale.value_or(
+            1.0 / std::sqrt(static_cast<double>(q_shape_[2])))),
+        false, causal, static_cast<int>(window_size[0]),
+        static_cast<int>(window_size[1]), static_cast<float>(softcap),
+        return_attn_probs && dropout_p > 0.0, generator);
+    assert(!result.empty() &&
+           "Linked `flash_attn_varlen_func` provider returned no output.");
+    at_out.copy_(result[0]);
+    if (at_softmax_lse.has_value()) {
+      assert(result.size() >= 3 &&
+             "Linked `flash_attn_varlen_func` provider did not return "
+             "auxiliary outputs.");
+      at_softmax_lse->copy_(result[1]);
+      at_s_dmask->copy_(result[2]);
+    }
+
+    (void)deterministic;
+  }
+};
+
+}  // namespace infini::ops::linked::torch
+
+#endif  // INFINI_OPS_LINKED_TORCH_OPS_FLASH_ATTN_VARLEN_FUNC_H_

src/torch/ops/flash_attn_varlen_func/flash_attn_varlen_func.cc

@@ -21,11 +21,16 @@ void Operator<FlashAttnVarlenFunc, Device::Type::kNvidia, 8>::operator()(
     const std::vector<int64_t> window_size, const double softcap,
     const bool deterministic, const bool return_attn_probs, Tensor out,
     std::optional<Tensor> softmax_lse, std::optional<Tensor> s_dmask) const {
+  assert(!alibi_slopes.has_value() &&
+         "The PyTorch `FlashAttnVarlenFunc` provider does not support "
+         "`alibi_slopes`.");
+  assert(!block_table.has_value() &&
+         "The PyTorch `FlashAttnVarlenFunc` provider does not support "
+         "`block_table`.");
+
   (void)softcap;
-  (void)alibi_slopes;
   (void)deterministic;
   (void)return_attn_probs;
-  (void)block_table;
 
   const auto device_index = static_cast<c10::DeviceIndex>(device_index_);
   const c10::cuda::CUDAGuard device_guard{device_index};