flash_4_hub does not forward return_lse to its kernel #15180

dispatch_attention_fn accepts return_lse through attention_kwargs, but the flash_4_hub adapter does not pass it to flash_attn_func. The Hub kernel's return_lse default is False. For inference inputs that do not require gradients, the kernel therefore returns None for LSE even when the caller requested it.

The following reproduces the adapter behavior without a GPU or a Hub download. The stub implements the relevant return contract of kernels-community/flash-attn4:

import torch
from diffusers.models.attention_dispatch import (
    AttentionBackendName,
    _AttentionBackendRegistry,
    _HUB_KERNELS_REGISTRY,
    dispatch_attention_fn,
)

def flash_attn_func(*, q, k, v, softmax_scale, causal, return_lse=False):
    return q, torch.ones(1, 1, 2) if return_lse else None

_HUB_KERNELS_REGISTRY[AttentionBackendName.FLASH_4_HUB].kernel_fn = flash_attn_func
_AttentionBackendRegistry._checks_enabled = False
q = torch.zeros(1, 2, 1, 4)
out, lse = dispatch_attention_fn(
    q, q, q,
    backend=AttentionBackendName.FLASH_4_HUB,
    attention_kwargs={"return_lse": True},
)
print(lse)  # None on the base revision

Expected: lse is the tensor returned by the kernel when return_lse=True. Actual: None because the adapter leaves the kernel argument at its default.

Tested against diffusers main b674fd8c2c6890c0ae9856d3b73ef4fee309e7a1 with PyTorch 2.13.0. This reproduction uses a test double; a GPU kernel run has not been performed.