diff --git a/src/base/flash_attn_varlen_func.h b/src/base/flash_attn_varlen_func.h index ae911fd32..0fdf7ee95 100644 --- a/src/base/flash_attn_varlen_func.h +++ b/src/base/flash_attn_varlen_func.h @@ -34,7 +34,9 @@ class FlashAttnVarlenFunc : public Operator { false, false, std::nullopt, - out} {} + out, + std::nullopt, + std::nullopt} {} FlashAttnVarlenFunc(const Tensor q, const Tensor k, const Tensor v, const Tensor cu_seqlens_q, const Tensor cu_seqlens_k, @@ -45,25 +47,39 @@ class FlashAttnVarlenFunc : public Operator { const double softcap, const std::optional alibi_slopes, const bool deterministic, const bool return_attn_probs, - const std::optional block_table, Tensor out) + const std::optional block_table, Tensor out, + std::optional softmax_lse, + std::optional s_dmask) : q_shape_{q.shape()}, k_shape_{k.shape()}, v_shape_{v.shape()}, cu_seqlens_q_shape_{cu_seqlens_q.shape()}, cu_seqlens_k_shape_{cu_seqlens_k.shape()}, out_shape_{out.shape()}, + softmax_lse_shape_{softmax_lse.has_value() ? softmax_lse->shape() + : Tensor::Shape{}}, + s_dmask_shape_{s_dmask.has_value() ? s_dmask->shape() + : Tensor::Shape{}}, q_strides_{q.strides()}, k_strides_{k.strides()}, v_strides_{v.strides()}, cu_seqlens_q_strides_{cu_seqlens_q.strides()}, cu_seqlens_k_strides_{cu_seqlens_k.strides()}, out_strides_{out.strides()}, + softmax_lse_strides_{softmax_lse.has_value() ? softmax_lse->strides() + : Tensor::Strides{}}, + s_dmask_strides_{s_dmask.has_value() ? s_dmask->strides() + : Tensor::Strides{}}, q_dtype_{q.dtype()}, k_dtype_{k.dtype()}, v_dtype_{v.dtype()}, cu_seqlens_q_dtype_{cu_seqlens_q.dtype()}, cu_seqlens_k_dtype_{cu_seqlens_k.dtype()}, 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"); @@ -77,6 +93,22 @@ class FlashAttnVarlenFunc : public Operator { "and no greater than 256"); assert(out.shape() == q.shape() && "`FlashAttnVarlenFunc` output must have the same shape as Q"); + assert(softmax_lse.has_value() == s_dmask.has_value() && + "`FlashAttnVarlenFunc` auxiliary outputs must be provided " + "together"); + assert(return_attn_probs == has_auxiliary_outputs_ && + "`FlashAttnVarlenFunc` requires auxiliary outputs exactly when " + "`return_attn_probs` is true"); + if (has_auxiliary_outputs_) { + assert((softmax_lse->shape() == Tensor::Shape{q.size(1), q.size(0)} && + softmax_lse_dtype_ == DataType::kFloat32 && + "`FlashAttnVarlenFunc` softmax LSE output must have shape " + "(num_heads, total_q) and dtype float32")); + assert(s_dmask->shape() == Tensor::Shape{0} && + s_dmask_dtype_ == q_dtype_ && + "`FlashAttnVarlenFunc` inference attention mask output must be " + "empty and match the Q dtype"); + } assert( (q_dtype_ == DataType::kFloat16 || q_dtype_ == DataType::kBFloat16) && q_dtype_ == k_dtype_ && q_dtype_ == v_dtype_ && @@ -108,8 +140,6 @@ class FlashAttnVarlenFunc : public Operator { "`FlashAttnVarlenFunc` does not yet support softcap"); assert(!deterministic && "`FlashAttnVarlenFunc` does not yet support deterministic mode"); - assert(!return_attn_probs && - "`FlashAttnVarlenFunc` does not yet return attention probabilities"); assert(!block_table.has_value() && "`FlashAttnVarlenFunc` does not yet support paged KV cache"); assert(!alibi_slopes.has_value() && @@ -122,6 +152,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) && + (!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"); (void)softmax_scale; @@ -134,7 +166,7 @@ class FlashAttnVarlenFunc : public Operator { Tensor out) const { (*this)(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, 0.0, std::nullopt, false, {-1, -1}, 0.0, std::nullopt, false, false, - std::nullopt, out); + std::nullopt, out, std::nullopt, std::nullopt); } virtual void operator()( @@ -145,7 +177,8 @@ class FlashAttnVarlenFunc : public Operator { const std::vector window_size, const double softcap, const std::optional alibi_slopes, const bool deterministic, const bool return_attn_probs, const std::optional block_table, - Tensor out) const = 0; + Tensor out, std::optional softmax_lse, + std::optional s_dmask) const = 0; protected: Tensor::Shape q_shape_; @@ -160,6 +193,10 @@ class FlashAttnVarlenFunc : public Operator { Tensor::Shape out_shape_; + Tensor::Shape softmax_lse_shape_; + + Tensor::Shape s_dmask_shape_; + Tensor::Strides q_strides_; Tensor::Strides k_strides_; @@ -172,6 +209,10 @@ class FlashAttnVarlenFunc : public Operator { Tensor::Strides out_strides_; + Tensor::Strides softmax_lse_strides_; + + Tensor::Strides s_dmask_strides_; + DataType q_dtype_; DataType k_dtype_; @@ -184,6 +225,12 @@ class FlashAttnVarlenFunc : public Operator { DataType out_dtype_; + DataType softmax_lse_dtype_; + + DataType s_dmask_dtype_; + + bool has_auxiliary_outputs_{false}; + int device_index_{0}; }; 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 2569c28ca..ab3a350f5 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 @@ -19,7 +19,8 @@ void Operator::operator()( const std::vector window_size, const double softcap, const std::optional alibi_slopes, const bool deterministic, const bool return_attn_probs, const std::optional block_table, - Tensor out) const { + Tensor out, std::optional softmax_lse, + std::optional s_dmask) const { (void)softcap; (void)alibi_slopes; (void)deterministic; @@ -47,6 +48,16 @@ void Operator::operator()( 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_softmax_lse; + std::optional at_s_dmask; + 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_)); + } const std::optional window_size_left = window_size[0] < 0 ? std::nullopt @@ -61,9 +72,13 @@ void Operator::operator()( max_seqlen_k, dropout_p, causal, false, softmax_scale, window_size_left, window_size_right, std::nullopt, std::nullopt); - // ATen owns the returned tensor. Keep the InfiniOps trailing-output ABI by - // copying it into the caller-provided buffer on the selected CUDA stream. + // ATen owns the returned tensors. Keep the InfiniOps trailing-output ABI + // by copying them into caller-provided buffers on the selected CUDA stream. at_out.copy_(std::get<0>(result)); + if (at_softmax_lse.has_value()) { + at_softmax_lse->copy_(std::get<1>(result)); + at_s_dmask->copy_(std::get<4>(result)); + } }; const c10::cuda::CUDAStreamGuard stream_guard{ diff --git a/src/torch/ops/flash_attn_varlen_func/flash_attn_varlen_func.h b/src/torch/ops/flash_attn_varlen_func/flash_attn_varlen_func.h index faf07ce46..b2e16d182 100644 --- a/src/torch/ops/flash_attn_varlen_func/flash_attn_varlen_func.h +++ b/src/torch/ops/flash_attn_varlen_func/flash_attn_varlen_func.h @@ -20,8 +20,9 @@ class Operator const std::vector window_size, const double softcap, const std::optional alibi_slopes, const bool deterministic, const bool return_attn_probs, - const std::optional block_table, - Tensor out) const override; + const std::optional block_table, Tensor out, + std::optional softmax_lse, + std::optional s_dmask) const override; }; } // namespace infini::ops diff --git a/tests/test_flash_attn_varlen_func.py b/tests/test_flash_attn_varlen_func.py index 1febffdcc..80e778e35 100644 --- a/tests/test_flash_attn_varlen_func.py +++ b/tests/test_flash_attn_varlen_func.py @@ -55,6 +55,12 @@ def test_flash_attn_varlen_func( cu_seqlens_q = _cumulative_lengths(q_lens, device) cu_seqlens_k = _cumulative_lengths(k_lens, device) out = torch.empty_like(q) + softmax_lse = torch.empty( + (q.size(1), q.size(0)), + dtype=torch.float32, + device=q.device, + ) + s_dmask = torch.empty((0,), dtype=q.dtype, device=q.device) infini.ops.flash_attn_varlen_func( q, @@ -71,9 +77,11 @@ def test_flash_attn_varlen_func( 0.0, None, False, - False, + True, None, out, + softmax_lse, + s_dmask, stream=get_stream(q.device), implementation_index=implementation_index, ) @@ -89,6 +97,25 @@ def test_flash_attn_varlen_func( window_size, ) 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] + ), + ) + torch.testing.assert_close(softmax_lse, expected_auxiliary[1]) + torch.testing.assert_close(s_dmask, expected_auxiliary[4]) def test_flash_attn_varlen_func_non_default_stream(device, implementation_index): @@ -125,6 +152,8 @@ def test_flash_attn_varlen_func_non_default_stream(device, implementation_index) False, None, out, + None, + None, stream=stream.cuda_stream, implementation_index=implementation_index, )