Skip to content

feat!: return flash_attn_varlen_func auxiliary outputs #854

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
voltjia merged 1 commit into from
Jul 31, 2026
Merged

feat!: return flash_attn_varlen_func auxiliary outputs #854

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
59 changes: 53 additions & 6 deletions src/base/flash_attn_varlen_func.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,9 @@ class FlashAttnVarlenFunc : public Operator<FlashAttnVarlenFunc> {
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,
Expand All @@ -45,25 +47,39 @@ class FlashAttnVarlenFunc : public Operator<FlashAttnVarlenFunc> {
const double softcap,
const std::optional<Tensor> alibi_slopes,
const bool deterministic, const bool return_attn_probs,
const std::optional<Tensor> block_table, Tensor out)
const std::optional<Tensor> block_table, Tensor out,
std::optional<Tensor> softmax_lse,
std::optional<Tensor> 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");
Expand All @@ -77,6 +93,22 @@ class FlashAttnVarlenFunc : public Operator<FlashAttnVarlenFunc> {
"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_ &&
Expand Down Expand Up @@ -108,8 +140,6 @@ class FlashAttnVarlenFunc : public Operator<FlashAttnVarlenFunc> {
"`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() &&
Expand All @@ -122,6 +152,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) &&
(!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;
Expand All @@ -134,7 +166,7 @@ class FlashAttnVarlenFunc : public Operator<FlashAttnVarlenFunc> {
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()(
Expand All @@ -145,7 +177,8 @@ class FlashAttnVarlenFunc : public Operator<FlashAttnVarlenFunc> {
const std::vector<int64_t> window_size, const double softcap,
const std::optional<Tensor> alibi_slopes, const bool deterministic,
const bool return_attn_probs, const std::optional<Tensor> block_table,
Tensor out) const = 0;
Tensor out, std::optional<Tensor> softmax_lse,
std::optional<Tensor> s_dmask) const = 0;

protected:
Tensor::Shape q_shape_;
Expand All @@ -160,6 +193,10 @@ class FlashAttnVarlenFunc : public Operator<FlashAttnVarlenFunc> {

Tensor::Shape out_shape_;

Tensor::Shape softmax_lse_shape_;

Tensor::Shape s_dmask_shape_;

Tensor::Strides q_strides_;

Tensor::Strides k_strides_;
Expand All @@ -172,6 +209,10 @@ class FlashAttnVarlenFunc : public Operator<FlashAttnVarlenFunc> {

Tensor::Strides out_strides_;

Tensor::Strides softmax_lse_strides_;

Tensor::Strides s_dmask_strides_;

DataType q_dtype_;

DataType k_dtype_;
Expand All @@ -184,6 +225,12 @@ class FlashAttnVarlenFunc : public Operator<FlashAttnVarlenFunc> {

DataType out_dtype_;

DataType softmax_lse_dtype_;

DataType s_dmask_dtype_;

bool has_auxiliary_outputs_{false};

int device_index_{0};
};

Expand Down
21 changes: 18 additions & 3 deletions src/torch/ops/flash_attn_varlen_func/flash_attn_varlen_func.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,8 @@ void Operator<FlashAttnVarlenFunc, Device::Type::kNvidia, 8>::operator()(
const std::vector<int64_t> window_size, const double softcap,
const std::optional<Tensor> alibi_slopes, const bool deterministic,
const bool return_attn_probs, const std::optional<Tensor> block_table,
Tensor out) const {
Tensor out, std::optional<Tensor> softmax_lse,
std::optional<Tensor> s_dmask) const {
(void)softcap;
(void)alibi_slopes;
(void)deterministic;
Expand Down Expand Up @@ -47,6 +48,16 @@ void Operator<FlashAttnVarlenFunc, Device::Type::kNvidia, 8>::operator()(
cu_seqlens_k_strides_, cu_seqlens_k_dtype_, device_index_);
auto at_out = ToAtenTensor<Device::Type::kNvidia>(
out.data(), out_shape_, out_strides_, out_dtype_, device_index_);
std::optional<at::Tensor> at_softmax_lse;
std::optional<at::Tensor> at_s_dmask;
if (softmax_lse.has_value()) {
at_softmax_lse.emplace(ToAtenTensor<Device::Type::kNvidia>(
softmax_lse->data(), softmax_lse_shape_, softmax_lse_strides_,
softmax_lse_dtype_, device_index_));
at_s_dmask.emplace(ToAtenTensor<Device::Type::kNvidia>(
s_dmask->data(), s_dmask_shape_, s_dmask_strides_, s_dmask_dtype_,
device_index_));
}

const std::optional<int64_t> window_size_left =
window_size[0] < 0 ? std::nullopt
Expand All @@ -61,9 +72,13 @@ void Operator<FlashAttnVarlenFunc, Device::Type::kNvidia, 8>::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{
Expand Down
5 changes: 3 additions & 2 deletions src/torch/ops/flash_attn_varlen_func/flash_attn_varlen_func.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,8 +20,9 @@ class Operator<FlashAttnVarlenFunc, Device::Type::kNvidia, 8>
const std::vector<int64_t> window_size, const double softcap,
const std::optional<Tensor> alibi_slopes,
const bool deterministic, const bool return_attn_probs,
const std::optional<Tensor> block_table,
Tensor out) const override;
const std::optional<Tensor> block_table, Tensor out,
std::optional<Tensor> softmax_lse,
std::optional<Tensor> s_dmask) const override;
};

} // namespace infini::ops
Expand Down
31 changes: 30 additions & 1 deletion tests/test_flash_attn_varlen_func.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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,
)
Expand All @@ -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):
Expand Down Expand Up @@ -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,
)
Expand Down
Loading