Repository: huggingface/diffusers
File: src/diffusers/models/attention_dispatch.py (Python)

Code before:

def _flash_attention_3_hub_backward_op(
    ctx: torch.autograd.function.FunctionCtx,
    grad_out: torch.Tensor,
    *args,
    window_size: tuple[int, int] = (-1, -1),
    softcap: float = 0.0,
    num_splits: int = 1,
    pack_gqa: bool | None = None,
    deterministic: bool = False,
    sm_margin: int = 0,
):
    query, key, value = ctx.saved_tensors
    kernel_fn = ctx._hub_kernel
    with torch.enable_grad():
        query_r = query.detach().requires_grad_(True)
        key_r = key.detach().requires_grad_(True)
        value_r = value.detach().requires_grad_(True)

        out = kernel_fn(
            q=query_r,
            k=key_r,
            v=value_r,
            softmax_scale=ctx.scale,
            causal=ctx.is_causal,
            qv=None,
            q_descale=None,
            k_descale=None,
            v_descale=None,
            window_size=window_size,
            softcap=softcap,
            num_splits=num_splits,
            pack_gqa=pack_gqa,
            deterministic=deterministic,
            sm_margin=sm_margin,
            return_attn_probs=False,
        )
        if isinstance(out, tuple):
            out = out[0]

        grad_query, grad_key, grad_value = torch.autograd.grad(
            out,
            (query_r, key_r, value_r),
            grad_out,
            retain_graph=False,
            allow_unused=False,
        )

    return grad_query, grad_key, grad_value



Review comment (bug):
This would result in a second in a forward pass during the backward op right?  Would it make sense to just raise an error here similar to sage attention?

Code after:

def _flash_attention_3_hub_backward_op(
    ctx: torch.autograd.function.FunctionCtx,
    grad_out: torch.Tensor,
    *args,
    window_size: tuple[int, int] = (-1, -1),
    softcap: float = 0.0,
    num_splits: int = 1,
    pack_gqa: bool | None = None,
    deterministic: bool = False,
    sm_margin: int = 0,
):
    query, key, value = ctx.saved_tensors
    kernel_fn = ctx._hub_kernel
    # NOTE: Unlike the FA2 hub kernel, the FA3 hub kernel does not expose separate wrapped forward/backward
    # primitives (no `wrapped_forward_attr`/`wrapped_backward_attr` in its `_HubKernelConfig`). We
    # therefore rerun the forward pass under `torch.enable_grad()` and differentiate through it with
    # `torch.autograd.grad()`. This is a second forward pass during backward; it can be avoided once
    # the FA3 hub exposes a dedicated fused backward kernel (analogous to `_wrapped_flash_attn_backward`
    # in the FA2 hub), at which point this can be refactored to match `_flash_attention_hub_backward_op`.
    with torch.enable_grad():
        query_r = query.detach().requires_grad_(True)
        key_r = key.detach().requires_grad_(True)
        value_r = value.detach().requires_grad_(True)

        out = kernel_fn(
            q=query_r,
            k=key_r,
            v=value_r,
            softmax_scale=ctx.scale,
            causal=ctx.is_causal,
            qv=None,
            q_descale=None,
            k_descale=None,
            v_descale=None,
            window_size=window_size,
            softcap=softcap,
            num_splits=num_splits,
            pack_gqa=pack_gqa,
            deterministic=deterministic,
            sm_margin=sm_margin,
            return_attn_probs=False,
        )
        if isinstance(out, tuple):
            out = out[0]

        grad_query, grad_key, grad_value = torch.autograd.grad(
            out,
            (query_r, key_r, value_r),
            grad_out,
            retain_graph=False,