flash_varlen raises when requesting log-sum-exp #15182

dispatch_attention_fn accepts return_lse through attention_kwargs. The flash_varlen adapter forwards it to FlashAttention as return_attn_probs, but then calls unflatten on the returned tuple. With return_lse=True, FlashAttention returns the output, LSE, and an optional attention-probability tensor, so the adapter raises AttributeError: 'tuple' object has no attribute 'unflatten'.

With FlashAttention 2 installed on a supported CUDA device:

import torch
from diffusers.models.attention_dispatch import AttentionBackendName, dispatch_attention_fn

q = torch.randn(2, 32, 2, 64, device="cuda", dtype=torch.float16)
out, lse = dispatch_attention_fn(
    q, q, q,
    backend=AttentionBackendName.FLASH_VARLEN,
    attention_kwargs={"return_lse": True},
)

Expected: the output and packed log-sum-exp tensor are returned. Actual: the tuple is treated as a tensor and raises before returning.

Reproduced against diffusers main b674fd8c2c6890c0ae9856d3b73ef4fee309e7a1 on GH200 with PyTorch 2.14.1 and a source-built FlashAttention 2.8.4 kernel. The kernel function was bound directly to the adapter because this environment did not have FlashAttention distribution metadata installed.