diff --git a/src/linked/torch/cambricon/ops/flash_attn_varlen_func/flash_attn.cc b/src/linked/torch/cambricon/ops/flash_attn_varlen_func/flash_attn.cc new file mode 100644 index 000000000..c3c6eceaf --- /dev/null +++ b/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 + +#include +#include + +#include "common/op_utils/paged_kv_cache.h" +#include "torch/cambricon/c10.h" +#include "torch/tensor_.h" + +std::vector mha_varlen_fwd( + const at::Tensor& q, const at::Tensor& k, const at::Tensor& v, + std::optional& out, const at::Tensor& cu_seqlens_q, + const at::Tensor& cu_seqlens_k, std::optional& seqused_k, + std::optional& 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 generator); + +namespace infini::ops { + +void Operator::operator()( + const Tensor q, const Tensor k, const Tensor v, const Tensor cu_seqlens_q, + const Tensor cu_seqlens_k, const std::optional alibi_slopes, + const std::optional block_table, const int64_t max_seqlen_q, + const int64_t max_seqlen_k, const double dropout_p, + const std::optional softmax_scale, const bool causal, + const std::vector window_size, const double softcap, + const bool deterministic, const bool return_attn_probs, Tensor out, + std::optional softmax_lse, std::optional s_dmask) const { + const C10::StreamGuard stream_guard{ + C10::GetStreamFromExternal(stream_, + device_index_)}; + auto at_q = ToAtenTensor( + const_cast(q.data()), q_shape_, q_strides_, q_dtype_, + device_index_); + auto at_k = ToAtenTensor( + const_cast(k.data()), k_shape_, k_strides_, k_dtype_, + device_index_); + auto at_v = ToAtenTensor( + const_cast(v.data()), v_shape_, v_strides_, v_dtype_, + device_index_); + auto at_cu_seqlens_q = ToAtenTensor( + const_cast(cu_seqlens_q.data()), cu_seqlens_q_shape_, + cu_seqlens_q_strides_, cu_seqlens_q_dtype_, device_index_); + auto at_cu_seqlens_k = ToAtenTensor( + const_cast(cu_seqlens_k.data()), cu_seqlens_k_shape_, + cu_seqlens_k_strides_, cu_seqlens_k_dtype_, device_index_); + auto at_out = ToAtenTensor( + out.data(), out_shape_, out_strides_, out_dtype_, device_index_); + + std::optional at_out_optional{at_out}; + std::optional at_seqused_k; + std::optional at_alibi_slopes; + std::optional generator; + if (alibi_slopes.has_value()) { + at_alibi_slopes.emplace(ToAtenTensor( + const_cast(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( + const_cast(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(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 packed_k; + std::vector 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(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(max_seqlen_q), static_cast(max_seqlen_k), + static_cast(dropout_p), + static_cast(softmax_scale.value_or( + 1.0 / std::sqrt(static_cast(q_shape_[2])))), + false, causal, static_cast(window_size[0]), + static_cast(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( + 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 diff --git a/src/linked/torch/cambricon/ops/flash_attn_varlen_func/flash_attn.h b/src/linked/torch/cambricon/ops/flash_attn_varlen_func/flash_attn.h new file mode 100644 index 000000000..2e8e4506e --- /dev/null +++ b/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 + : 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 alibi_slopes, + const std::optional block_table, + const int64_t max_seqlen_q, const int64_t max_seqlen_k, + const double dropout_p, + const std::optional softmax_scale, const bool causal, + const std::vector window_size, const double softcap, + const bool deterministic, const bool return_attn_probs, + Tensor out, std::optional softmax_lse, + std::optional s_dmask) const override; +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_LINKED_TORCH_CAMBRICON_OPS_FLASH_ATTN_VARLEN_FUNC_FLASH_ATTN_H_ diff --git a/src/linked/torch/cambricon/ops/flash_attn_varlen_func/flash_attn.yaml b/src/linked/torch/cambricon/ops/flash_attn_varlen_func/flash_attn.yaml new file mode 100644 index 000000000..f45ca5758 --- /dev/null +++ b/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 const&, at::Tensor const&, std::optional&, std::optional&, int, int, float, float, bool, bool, int, int, bool, std::optional) diff --git a/tests/test_flash_attn_varlen_func.py b/tests/test_flash_attn_varlen_func.py index 5bdf5d514..6e7050a6d 100644 --- a/tests/test_flash_attn_varlen_func.py +++ b/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