feat(cambricon): add varlen flash attention provider

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

@@ -0,0 +1,121 @@
+#include "linked/torch/cambricon/ops/flash_attn_varlen_func/flash_attn.h"
+
+#include <ATen/core/Generator.h>
+
+#include <cassert>
+#include <cmath>
+
+#include "common/op_utils/paged_kv_cache.h"
+#include "torch/cambricon/c10.h"
+#include "torch/tensor_.h"
+
+std::vector<at::Tensor> mha_varlen_fwd(
+    const 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<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, bool return_softmax,
+    std::optional<at::Generator> generator);
+
+namespace infini::ops {
+
+void Operator<FlashAttnVarlenFunc, Device::Type::kCambricon, 16>::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 {
+  const C10<Device::Type::kCambricon>::StreamGuard stream_guard{
+      C10<Device::Type::kCambricon>::GetStreamFromExternal(stream_,
+                                                           device_index_)};
+  auto at_q = ToAtenTensor<Device::Type::kCambricon>(
+      const_cast<void*>(q.data()), q_shape_, q_strides_, q_dtype_,
+      device_index_);
+  auto at_k = ToAtenTensor<Device::Type::kCambricon>(
+      const_cast<void*>(k.data()), k_shape_, k_strides_, k_dtype_,
+      device_index_);
+  auto at_v = ToAtenTensor<Device::Type::kCambricon>(
+      const_cast<void*>(v.data()), v_shape_, v_strides_, v_dtype_,
+      device_index_);
+  auto at_cu_seqlens_q = ToAtenTensor<Device::Type::kCambricon>(
+      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<Device::Type::kCambricon>(
+      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<Device::Type::kCambricon>(
+      out.data(), out_shape_, out_strides_, out_dtype_, device_index_);
+
+  std::optional<at::Tensor> at_out_optional{at_out};
+  std::optional<at::Tensor> at_seqused_k;
+  std::optional<at::Tensor> at_alibi_slopes;
+  std::optional<at::Generator> generator;
+  if (alibi_slopes.has_value()) {
+    at_alibi_slopes.emplace(ToAtenTensor<Device::Type::kCambricon>(
+        const_cast<void*>(alibi_slopes->data()), alibi_slopes_shape_,
+        alibi_slopes_strides_, alibi_slopes_dtype_, device_index_));
+  }
+
+  at::Tensor at_k_for_call = at_k;
+  at::Tensor at_v_for_call = at_v;
+  if (block_table.has_value()) {
+    auto at_block_table = ToAtenTensor<Device::Type::kCambricon>(
+        const_cast<void*>(block_table->data()), block_table_shape_,
+        block_table_strides_, block_table_dtype_, device_index_);
+    const auto host_cu_seqlens_k =
+        paged_kv_cache::ToHostInt32Vector(at_cu_seqlens_k);
+    const auto host_block_table =
+        paged_kv_cache::ToHostInt32Vector(at_block_table);
+    const int64_t batch_size =
+        static_cast<int64_t>(host_cu_seqlens_k.size()) - 1;
+    const int64_t table_width = at_block_table.size(1);
+    assert(at_block_table.size(0) == batch_size &&
+           "KV cache block table batch size does not match cu_seqlens_k");
+
+    std::vector<at::Tensor> packed_k;
+    std::vector<at::Tensor> packed_v;
+    packed_k.reserve(batch_size);
+    packed_v.reserve(batch_size);
+    for (int64_t batch = 0; batch < batch_size; ++batch) {
+      const int64_t length = static_cast<int64_t>(host_cu_seqlens_k[batch + 1] -
+                                                  host_cu_seqlens_k[batch]);
+      assert(length >= 0 && "cu_seqlens_k must be nondecreasing");
+      packed_k.push_back(paged_kv_cache::GatherSequence(
+          at_k, host_block_table, table_width, batch, length));
+      packed_v.push_back(paged_kv_cache::GatherSequence(
+          at_v, host_block_table, table_width, batch, length));
+    }
+    at_k_for_call = at::cat(packed_k, 0).contiguous();
+    at_v_for_call = at::cat(packed_v, 0).contiguous();
+  }
+
+  const auto result = ::mha_varlen_fwd(
+      at_q, at_k_for_call, at_v_for_call, at_out_optional, at_cu_seqlens_q,
+      at_cu_seqlens_k, at_seqused_k, 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]), false, generator);
+  assert(result.size() > 5 &&
+         "Cambricon FlashAttention returned an incomplete result");
+  at_out.copy_(result[0]);
+
+  if (return_attn_probs) {
+    auto at_softmax_lse = ToAtenTensor<Device::Type::kCambricon>(
+        softmax_lse->data(), softmax_lse_shape_, softmax_lse_strides_,
+        softmax_lse_dtype_, device_index_);
+    at_softmax_lse.copy_(result[5]);
+  }
+
+  (void)softcap;
+  (void)deterministic;
+  (void)s_dmask;
+}
+
+}  // namespace infini::ops

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

@@ -0,0 +1,30 @@
+#ifndef INFINI_OPS_LINKED_TORCH_CAMBRICON_OPS_FLASH_ATTN_VARLEN_FUNC_FLASH_ATTN_H_
+#define INFINI_OPS_LINKED_TORCH_CAMBRICON_OPS_FLASH_ATTN_VARLEN_FUNC_FLASH_ATTN_H_
+
+#include "base/flash_attn_varlen_func.h"
+
+namespace infini::ops {
+
+template <>
+class Operator<FlashAttnVarlenFunc, Device::Type::kCambricon, 16>
+    : public FlashAttnVarlenFunc {
+ public:
+  using FlashAttnVarlenFunc::FlashAttnVarlenFunc;
+  using 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;
+};
+
+}  // namespace infini::ops
+
+#endif  // INFINI_OPS_LINKED_TORCH_CAMBRICON_OPS_FLASH_ATTN_VARLEN_FUNC_FLASH_ATTN_H_

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

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

tests/test_flash_attn_varlen_func.py

@@ -53,8 +53,8 @@ def test_flash_attn_varlen_func(
     rtol,
     atol,
 ):
-    if device not in ("cuda", "musa"):
-        pytest.skip("FlashAttention requires the NVIDIA or Moore backend")
+    if device not in ("cuda", "musa", "mlu"):
+        pytest.skip("FlashAttention requires the NVIDIA, Moore, or Cambricon backend")
     if device == "musa" and window_size != (-1, -1):
         pytest.skip("TorchMusa FlashAttention does not support local windows")
     if device == "musa" and not paged and causal and q_lens != k_lens:
@@ -157,38 +157,57 @@ def test_flash_attn_varlen_func(
     torch.testing.assert_close(out, expected, rtol=rtol, atol=atol)
 
     if return_attn_probs:
-        reference_window_size_right = None
-        if device == "cuda" and causal:
-            reference_window_size_right = 0
-        elif device == "cuda" and window_size[1] >= 0:
-            reference_window_size_right = window_size[1]
-
-        expected_auxiliary = torch.ops.aten._flash_attention_forward.default(
-            q,
-            k,
-            v,
-            cu_seqlens_q,
-            cu_seqlens_k,
-            max(q_lens),
-            max(k_lens),
-            0.0,
-            causal,
-            False,
-            scale=scale,
-            window_size_left=None if window_size[0] < 0 else window_size[0],
-            window_size_right=reference_window_size_right,
-        )
-        expected_softmax_lse = _pack_varlen_softmax_lse(
-            expected_auxiliary[1],
-            q_lens,
-        )
+        if device == "mlu":
+            expected_softmax_lse = _reference_varlen_softmax_lse(
+                q,
+                k,
+                q_lens,
+                k_lens,
+                scale,
+                causal,
+                window_size,
+                alibi_slopes,
+            )
+        else:
+            reference_window_size_right = None
+            if device == "cuda" and causal:
+                reference_window_size_right = 0
+            elif device == "cuda" and window_size[1] >= 0:
+                reference_window_size_right = window_size[1]
+
+            expected_auxiliary = torch.ops.aten._flash_attention_forward.default(
+                q,
+                k,
+                v,
+                cu_seqlens_q,
+                cu_seqlens_k,
+                max(q_lens),
+                max(k_lens),
+                0.0,
+                causal,
+                False,
+                scale=scale,
+                window_size_left=None if window_size[0] < 0 else window_size[0],
+                window_size_right=reference_window_size_right,
+            )
+            expected_softmax_lse = _pack_varlen_softmax_lse(
+                expected_auxiliary[1],
+                q_lens,
+            )
         torch.testing.assert_close(softmax_lse, expected_softmax_lse)
-        torch.testing.assert_close(s_dmask, expected_auxiliary[4])
+        if device == "cuda":
+            torch.testing.assert_close(s_dmask, expected_auxiliary[4])
 
 
 def test_flash_attn_varlen_func_non_default_stream(device, implementation_index):
-    if device != "cuda":
-        pytest.skip("non-default CUDA streams require the NVIDIA backend")
+    if device == "cuda":
+        accelerator = torch.cuda
+        stream_attribute = "cuda_stream"
+    elif device == "mlu":
+        accelerator = torch.mlu
+        stream_attribute = "mlu_stream"
+    else:
+        pytest.skip("stream coverage requires an accelerator backend")
 
     dtype = torch.float16
     q_lens = (3, 5)
@@ -199,8 +218,8 @@ def test_flash_attn_varlen_func_non_default_stream(device, implementation_index)
     cu_seqlens_q = _cumulative_lengths(q_lens, device)
     cu_seqlens_k = _cumulative_lengths(k_lens, device)
     out = torch.empty_like(q)
-    stream = torch.cuda.Stream()
-    stream.wait_stream(torch.cuda.current_stream())
+    stream = accelerator.Stream()
+    stream.wait_stream(accelerator.current_stream())
 
     infini.ops.flash_attn_varlen_func(
         q,
@@ -222,7 +241,7 @@ def test_flash_attn_varlen_func_non_default_stream(device, implementation_index)
         out,
         None,
         None,
-        stream=stream.cuda_stream,
+        stream=getattr(stream, stream_attribute),
         implementation_index=implementation_index,
     )
 
@@ -283,8 +302,8 @@ def test_flash_attn_varlen_func_default_stream(device, implementation_index):
 
 
 def test_flash_attn_varlen_func_defaults(device, implementation_index):
-    if device not in ("cuda", "musa"):
-        pytest.skip("FlashAttention requires the NVIDIA or Moore backend")
+    if device not in ("cuda", "musa", "mlu"):
+        pytest.skip("FlashAttention requires the NVIDIA, Moore, or Cambricon backend")
 
     q = torch.randn((5, 4, 64), dtype=torch.float16, device=device)
     k = torch.randn((5, 4, 64), dtype=torch.float16, device=device)
@@ -629,6 +648,49 @@ def _reference_varlen_attention(
     return torch.cat(outputs)
 
 
+def _reference_varlen_softmax_lse(
+    q,
+    k,
+    q_lens,
+    k_lens,
+    scale,
+    causal,
+    window_size,
+    alibi_slopes=None,
+):
+    outputs = []
+    q_offset = 0
+    k_offset = 0
+
+    for batch_index, (q_len, k_len) in enumerate(zip(q_lens, k_lens)):
+        q_seq = q[q_offset : q_offset + q_len].transpose(0, 1)
+        k_seq = k[k_offset : k_offset + k_len].transpose(0, 1)
+        groups = q_seq.size(0) // k_seq.size(0)
+        k_seq = k_seq.repeat_interleave(groups, dim=0)
+        scale_factor = scale if scale is not None else 1.0 / math.sqrt(q.size(-1))
+        scores = (
+            torch.matmul(q_seq.float(), k_seq.float().transpose(-2, -1)) * scale_factor
+        )
+        if alibi_slopes is not None:
+            slopes = (
+                alibi_slopes if alibi_slopes.ndim == 1 else alibi_slopes[batch_index]
+            )
+            query_positions = torch.arange(q_len, device=q.device).unsqueeze(1)
+            key_positions = torch.arange(k_len, device=q.device).unsqueeze(0)
+            distance = (query_positions + k_len - q_len - key_positions).abs()
+            scores += -slopes[:, None, None] * distance
+
+        mask = _attention_mask(q_len, k_len, causal, window_size, q.device)
+        if mask is not None:
+            scores.masked_fill_(~mask.unsqueeze(0), -math.inf)
+        softmax_lse = torch.logsumexp(scores, dim=-1)
+        outputs.append(torch.where(torch.isneginf(softmax_lse), math.inf, softmax_lse))
+        q_offset += q_len
+        k_offset += k_len
+
+    return torch.cat(outputs, dim=1)
+
+
 def _attention_mask(q_len, k_len, causal, window_size, device):
     left, right = window_size