add epilogue hook to flex attention (#45841)
Signed-off-by: Angel Li <liangel@meta.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>
liangel-02 committed
Jul 28, 2026 at 07:26 UTC
9b9fc4039c25a6e4fe0ae97361b62edd74b8b47e
1 file changed
+9
vllm/v1/attention/backends/flex_attention.py
+9
@@ -12,6 +12,7 @@ import torch
12
import torch._dynamo.decorators
13
import torch.nn.functional as F
14
from torch.nn.attention.flex_attention import (
15
+ AuxRequest,
16
BlockMask,
17
_mask_mod_signature,
18
_score_mod_signature,
@@ -1228,6 +1229,9 @@ class FlexAttentionImpl(AttentionImpl):
1229
if block_n is not None:
1230
self.block_n = block_n
1231
1232
+ # Optional post-attention epilogue transform
1233
+ self.out_transform = kwargs.get("out_transform")
1234
+
1235
@staticmethod
1236
def view_as_4d(tensor: torch.Tensor) -> torch.Tensor:
1237
"""View a 3d tensor as 4D."""
@@ -1392,8 +1396,13 @@ class FlexAttentionImpl(AttentionImpl):
1396
self.scale,
1397
enable_gqa=enable_gqa,
1398
kernel_options=kernel_options,
1399
+ return_aux=AuxRequest(lse=True) if self.out_transform is not None else None,
1400
)
1401
1402
+ if self.out_transform is not None:
1403
+ out, aux = out
1404
+ out = self.out_transform(out, aux.lse)
1405
+
1406
# Flex doesn't have an out variant today, rely on epilogue fusion
1407
out = out.permute(0, 2, 1, 3).squeeze(0)
1408
output[:num_actual_tokens, :, :].copy_(out)