Skip to content

set_attention_backend on one model leaks globally to other models via the active-backend registry #14249

Description

@sayakpaul

Calling model.set_attention_backend(...) on a single model does not only configure that model — it also flips a global active backend (_AttentionBackendRegistry.set_active_backend(...)). Any other model in the same process whose per-module attention backend is left unset then silently picks up that global backend inside dispatch_attention_fn (the backend is None fallback path).

This is surprising: setting the backend on model A changes the attention behavior of unrelated model B. It's especially harmful when B calls attention with arguments the leaked backend doesn't support — e.g. a masked attention (attn_mask=...) with a FlashAttention-3 backend, which raises ValueError: attn_mask is not supported for flash-attn 3. Even though the user never asked FA3 to run in B.

Reproduction

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

# Model A opts into a backend; nothing is configured on B.
# set_attention_backend() sets the module's processor._attention_backend AND, as a side
# effect, flips the process-global active backend:
_AttentionBackendRegistry.set_active_backend(AttentionBackendName("_flash_3_hub"))

# Model B just calls dispatch with backend=None (its per-module backend is unset) and a mask:
q = k = v = torch.randn(1, 4, 64, 64, device="cuda", dtype=torch.bfloat16)
mask = torch.zeros(1, 1, 64, 64, device="cuda", dtype=torch.bfloat16)

out = dispatch_attention_fn(q, k, v, attn_mask=mask, backend=None)  # -> uses the leaked global backend
# ValueError: `attn_mask` is not supported for flash-attn 3.

Cc: @dg845 @DN6

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