Skip to content

FLEX attention silently ignores is_causal=True when no mask is supplied #15195

Open

Description

@davrens101

Describe the bug

The FLEX attention backend silently ignores is_causal=True when no attention mask is supplied. _native_flex_attention first matches attn_mask is None, leaving block_mask=None; its subsequent causal-mask branch is unreachable for this supported argument combination.

Reproduced on upstream d3d8af8 with PyTorch 2.15.0+cpu and Diffusers 0.42.0.dev0. No model download is needed.

Reproduction

import torch
from torch.nn import functional as F
from diffusers.models.attention_dispatch import dispatch_attention_fn

torch.manual_seed(2026)
q, k, v = [torch.randn(1, 8, 2, 16) for _ in range(3)]
actual = dispatch_attention_fn(q, k, v, is_causal=True, backend="flex")
expected = F.scaled_dot_product_attention(
    q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), is_causal=True
).transpose(1, 2)
print((actual - expected).abs().max().item())

Observed maximum error: 2.117769241333008. The FLEX result instead matches noncausal SDPA within 3.5762786865234375e-07. Expected: causal output matching causal SDPA within floating-point tolerance.

Proposed scope and coordination

Build the causal block mask before accepting a missing mask as the noncausal case; preserve the existing explicit-mask validation. A focused fix is already proposed in #15188, which was opened without the required coordination issue. This report supplies the missing reproducible tracking; it does not claim prior maintainer acknowledgment. Please confirm whether this scope is appropriate before further contribution work proceeds.

AI assistance was used to reproduce and investigate this report. Verification here is CPU eager execution, not compiled CUDA coverage.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions