diff --git a/src/base/flash_attn_varlen_func.h b/src/base/flash_attn_varlen_func.h index dda248450..35ebb5aad 100644 --- a/src/base/flash_attn_varlen_func.h +++ b/src/base/flash_attn_varlen_func.h @@ -52,6 +52,12 @@ class FlashAttnVarlenFunc : public Operator { 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 { 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 { 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` 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 { 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 { 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 { 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 { DataType cu_seqlens_k_dtype_; + DataType alibi_slopes_dtype_; + + DataType block_table_dtype_; + DataType out_dtype_; DataType softmax_lse_dtype_; diff --git a/src/linked/torch/nvidia/ops/flash_attn_varlen_func/flash_attn.cc b/src/linked/torch/nvidia/ops/flash_attn_varlen_func/flash_attn.cc new file mode 100644 index 000000000..2b16d3678 --- /dev/null +++ b/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 mha_varlen_fwd( + 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& leftpad_k, + std::optional& block_table, + 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, float softcap, + bool return_softmax, std::optional generator); + +} // namespace flash + +namespace infini::ops::linked::torch::nvidia { + +std::vector FlashAttnVarlen::Call( + 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& leftpad_k, + std::optional& block_table, + 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, float softcap, + bool return_softmax, std::optional 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 diff --git a/src/linked/torch/nvidia/ops/flash_attn_varlen_func/flash_attn.h b/src/linked/torch/nvidia/ops/flash_attn_varlen_func/flash_attn.h new file mode 100644 index 000000000..9a02c7179 --- /dev/null +++ b/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 + +#include "linked/torch/nvidia/c10.h" +#include "linked/torch/ops/flash_attn_varlen_func.h" + +namespace infini::ops::linked::torch::nvidia { + +struct FlashAttnVarlen : C10 { + static std::vector Call( + 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& leftpad_k, + std::optional& block_table, + 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, float softcap, + bool return_softmax, std::optional 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 + : 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_ diff --git a/src/linked/torch/nvidia/ops/flash_attn_varlen_func/flash_attn.yaml b/src/linked/torch/nvidia/ops/flash_attn_varlen_func/flash_attn.yaml new file mode 100644 index 000000000..7f37cee62 --- /dev/null +++ b/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 const&, at::Tensor const&, std::optional&, std::optional&, std::optional&, std::optional&, int, int, float, float, bool, bool, int, int, float, bool, std::optional) diff --git a/src/linked/torch/ops/flash_attn_varlen_func.h b/src/linked/torch/ops/flash_attn_varlen_func.h new file mode 100644 index 000000000..2c7301fa3 --- /dev/null +++ b/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 + +#include +#include +#include +#include + +#include "base/flash_attn_varlen_func.h" +#include "torch/tensor_.h" + +namespace infini::ops::linked::torch { + +template +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 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 { + const typename Backend::StreamGuard stream_guard{ + Backend::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_alibi_slopes; + std::optional at_block_table; + std::optional at_softmax_lse; + std::optional at_s_dmask; + + 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_)); + } + if (block_table.has_value()) { + at_block_table.emplace(ToAtenTensor( + const_cast(block_table->data()), block_table_shape_, + block_table_strides_, block_table_dtype_, device_index_)); + } + if (softmax_lse.has_value()) { + at_softmax_lse.emplace(ToAtenTensor( + softmax_lse->data(), softmax_lse_shape_, softmax_lse_strides_, + softmax_lse_dtype_, device_index_)); + at_s_dmask.emplace(ToAtenTensor( + s_dmask->data(), s_dmask_shape_, s_dmask_strides_, s_dmask_dtype_, + device_index_)); + } + + std::optional at_out_optional; + std::optional at_seqused_k; + std::optional at_leftpad_k; + std::optional 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(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]), static_cast(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_ diff --git a/src/torch/ops/flash_attn_varlen_func/flash_attn_varlen_func.cc b/src/torch/ops/flash_attn_varlen_func/flash_attn_varlen_func.cc index c42fce589..a93de1cd7 100644 --- a/src/torch/ops/flash_attn_varlen_func/flash_attn_varlen_func.cc +++ b/src/torch/ops/flash_attn_varlen_func/flash_attn_varlen_func.cc @@ -21,11 +21,16 @@ void Operator::operator()( 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 { + 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(device_index_); const c10::cuda::CUDAGuard device_guard{device_index}; diff --git a/tests/test_flash_attn_varlen_func.py b/tests/test_flash_attn_varlen_func.py index 4d3b11aad..b045017cb 100644 --- a/tests/test_flash_attn_varlen_func.py +++ b/tests/test_flash_attn_varlen_func.py @@ -15,12 +15,16 @@ @pytest.mark.parametrize( - "q_lens, k_lens, num_heads, num_kv_heads, causal, window_size, scale", ( - ((3, 5), (4, 5), 4, 4, False, (-1, -1), None), - ((5, 2), (3, 6), 4, 2, True, (-1, -1), 0.125), - ((4, 3), (6, 2), 4, 2, False, (2, 1), None), - ((4, 3), (6, 2), 4, 2, True, (2, 1), None), + "q_lens, k_lens, num_heads, num_kv_heads, causal, window_size, " + "scale, paged, use_alibi" + ), + ( + ((3, 5), (4, 5), 4, 4, False, (-1, -1), None, False, False), + ((5, 2), (3, 6), 4, 2, True, (-1, -1), 0.125, False, False), + ((4, 3), (6, 2), 4, 2, False, (2, 1), None, False, False), + ((4, 3), (6, 2), 4, 2, True, (2, 1), None, False, False), + ((2, 3), (130, 300), 4, 2, True, (-1, -1), None, True, True), ), ) @pytest.mark.parametrize("head_dim", (64, 128)) @@ -39,6 +43,8 @@ def test_flash_attn_varlen_func( causal, window_size, scale, + paged, + use_alibi, head_dim, dtype, device, @@ -49,18 +55,63 @@ def test_flash_attn_varlen_func( if device != "cuda": pytest.skip("FlashAttention FA2 requires the NVIDIA backend") + if (paged or use_alibi) and implementation_index == 8: + pytest.skip("paged KV cache and ALiBi require the linked provider") + q = torch.randn((sum(q_lens), num_heads, head_dim), dtype=dtype, device=device) - k = torch.randn((sum(k_lens), num_kv_heads, head_dim), dtype=dtype, device=device) + block_table = None + if paged: + page_size = 256 + max_blocks = max((length + page_size - 1) // page_size for length in k_lens) + block_rows = [] + num_blocks = 0 + + for length in k_lens: + blocks = (length + page_size - 1) // page_size + row = list(range(num_blocks, num_blocks + blocks)) + row.extend([0] * (max_blocks - blocks)) + block_rows.append(row) + num_blocks += blocks + + k = torch.randn( + (num_blocks, page_size, num_kv_heads, head_dim), + dtype=dtype, + device=device, + ) + block_table = torch.tensor( + block_rows, + dtype=torch.int32, + device=device, + ) + else: + k = torch.randn( + (sum(k_lens), num_kv_heads, head_dim), + dtype=dtype, + device=device, + ) + v = torch.randn_like(k) cu_seqlens_q = _cumulative_lengths(q_lens, device) cu_seqlens_k = _cumulative_lengths(k_lens, device) + alibi_slopes = ( + torch.linspace(0.01, 0.04, num_heads, dtype=torch.float32, device=device) + if use_alibi + else None + ) out = torch.empty_like(q) - softmax_lse = torch.empty( - (q.size(1), q.size(0)), - dtype=torch.float32, - device=q.device, + return_attn_probs = not paged + softmax_lse = ( + torch.empty( + (q.size(1), q.size(0)), + dtype=torch.float32, + device=q.device, + ) + if return_attn_probs + else None + ) + s_dmask = ( + torch.empty((0,), dtype=q.dtype, device=q.device) if return_attn_probs else None ) - s_dmask = torch.empty((0,), dtype=q.dtype, device=q.device) infini.ops.flash_attn_varlen_func( q, @@ -68,8 +119,8 @@ def test_flash_attn_varlen_func( v, cu_seqlens_q, cu_seqlens_k, - None, - None, + alibi_slopes, + block_table, max(q_lens), max(k_lens), 0.0, @@ -78,7 +129,7 @@ def test_flash_attn_varlen_func( window_size, 0.0, False, - True, + return_attn_probs, out, softmax_lse, s_dmask, @@ -95,31 +146,35 @@ def test_flash_attn_varlen_func( scale, causal, window_size, + block_table, + alibi_slopes, ) torch.testing.assert_close(out, expected, rtol=rtol, atol=atol) - 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=( - 0 if causal else None if window_size[1] < 0 else window_size[1] - ), - ) - 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 return_attn_probs: + 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=( + 0 if causal else None if window_size[1] < 0 else window_size[1] + ), + ) + 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]) def test_flash_attn_varlen_func_non_default_stream(device, implementation_index): @@ -323,15 +378,26 @@ def _reference_varlen_attention( scale, causal, window_size, + block_table=None, + alibi_slopes=None, ): outputs = [] q_offset = 0 k_offset = 0 - for q_len, k_len in zip(q_lens, k_lens): + 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) - v_seq = v[k_offset : k_offset + k_len].transpose(0, 1) + if block_table is None: + k_seq = k[k_offset : k_offset + k_len] + v_seq = v[k_offset : k_offset + k_len] + else: + blocks = (k_len + k.size(1) - 1) // k.size(1) + block_indices = block_table[batch_index, :blocks].tolist() + k_seq = torch.cat(tuple(k[index] for index in block_indices))[:k_len] + v_seq = torch.cat(tuple(v[index] for index in block_indices))[:k_len] + + k_seq = k_seq.transpose(0, 1) + v_seq = v_seq.transpose(0, 1) groups = q_seq.size(0) // k_seq.size(0) k_seq = k_seq.repeat_interleave(groups, dim=0) v_seq = v_seq.repeat_interleave(groups, dim=0) @@ -346,6 +412,15 @@ def _reference_varlen_attention( 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 + if mask is not None: scores.masked_fill_(~mask.unsqueeze(0), -math.inf) probabilities = torch.softmax(scores, dim=-1)