Why `is_causal` is not used in `flash_attention_forward` ?

[lucaswychan]

I want to perform bidirectional attention in the Qwen3 model to train an embedding model, so I passed is_causal=False in the model forward (I manually added is_causal arguments in all forward method such as Qwen3Model and Qwen3Attention inmodeling_qwen3.py):

class Qwen3Attention(nn.Module):
    """Multi-headed attention from 'Attention Is All You Need' paper"""
        ...

    def forward(
        self,
        hidden_states: torch.Tensor,
        position_embeddings: tuple[torch.Tensor, torch.Tensor],
        attention_mask: Optional[torch.Tensor],
        past_key_value: Optional[Cache] = None,
        cache_position: Optional[torch.LongTensor] = None,
        is_causal: Optional[bool] = True,   # I add is_causal here
        **kwargs: Unpack[FlashAttentionKwargs],
    ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]:
        ...

        attn_output, attn_weights = attention_interface(
            self,
            query_states,
            key_states,
            value_states,
            attention_mask,
            dropout=0.0 if not self.training else self.attention_dropout,
            scaling=self.scaling,
            sliding_window=self.sliding_window,  # diff with Llama
            is_causal=is_causal, # and is_causal from the argument is passed to the attention_interface (e.g. `flash_attention_2`, `sdpa_attention_forward`)
            **kwargs,
        )

I can successfully change the causality of the attention in sdpa_attention_forward. However, I realized that it does not change the causality in the attention in flash_attention_forward. After diving into the implementation of flash_attention_forward, I found the reason in flash_attention_forward located at transformers/integrations/flash_attention.py:

def flash_attention_forward(
    module: torch.nn.Module,
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    attention_mask: Optional[torch.Tensor],
    dropout: float = 0.0,
    scaling: Optional[float] = None,
    sliding_window: Optional[int] = None,
    softcap: Optional[float] = None,
    **kwargs,
) -> tuple[torch.Tensor, None]:
    ...

    # FA2 always relies on the value set in the module, so remove it if present in kwargs to avoid passing it twice
    kwargs.pop("is_causal", None)

    attn_output = _flash_attention_forward(
        query,
        key,
        value,
        attention_mask,
        query_length=seq_len,
        is_causal=module.is_causal,  # here module is `Qwen3Attention`
        dropout=dropout,
        softmax_scale=scaling,
        sliding_window=sliding_window,
        softcap=softcap,
        use_top_left_mask=_use_top_left_mask,
        target_dtype=target_dtype,
        attn_implementation=module.config._attn_implementation,
        **kwargs,
    )

As you can see, the is_causal argument is popped, and the is_causal of Qwen3Attention is used as the argument. Note that Qwen3Attention.is_causal is never changed, and its default value is True, so the is_causal argument passed into _flash_attention_forward will always be True regardless of any change.

After I add a line of code to alter the Qwen3Attention.is_causal, i.e. self.is_causal = is_causal before passing the arguments into attention_interface, I can change the causality of flash_attention_forward. So I would like to know if it is a feature or a bug? Thank you!!

[Rocketknight1]

cc @ArthurZucker @Cyrilvallez @vasqu for FA2

[ArthurZucker]

More a feature than a bug! We should have flexibility for your case tho! Do you want to try and open a PR?

[lucaswychan]

@ArthurZucker I would love to open a PR and modify it. May I know the intention of popping is_causal in kwargs and using module.is_causal ? From my POV, it is more flexible to directly use is_causal argument instead of module.is_causal. Thank you !!!

[ArthurZucker]

well, it's easier to use is_causal from the module than having to always pass it maybe? @vasqu might remember better than me!

[vasqu]

Iirc, it's mostly a convention thing and most models refer to the module's is_causal attribute as it's easier for modeling code (as the model shouldn't change their attention behavior).

But I think allowing it as kwarg for power users would also be nice. However, this will likely need adjustments on

1. Other attention implementations as well e.g. sdpa
2. Mask creation (e.g. for sdpa)

The question is how do you intend to use this feature - if it's just reusing the attention module it shouldn't be an issue but as an end-to-end model it will require the mask creation point for example.

[lucaswychan]

It does convenient to relie on the attention module's is_causal, but since the attention module's is_causal is default as True and it is never changed, we are unable to change the module's attention behavior at all by relying on it.

TBH I manually added is_causal argument in all modules' forward (e.g. Qwen3Attention, Qwen3Model, Qwen3ModelForCausalLM). Besides, all of the default values of is_causal are True, except Qwen3Model. In this case, we can directly enable bidirectional attention without passing or setting any is_causal value. For instance:

import torch
from transformers import AutoTokenizer, AutoModel, AutoModelForCausalLM

model_name = "Qwen/Qwen3-4B"
device = torch.device("cuda")

tokenizer = AutoTokenizer.from_pretrained(model_name)


model = AutoModel.from_pretrained(model_name, device_map=device, attn_implementation="flash_attention_2", torch_dtype="auto").
# it uses the module `Qwen3Model`, with default value of `is_causal` as `False` 
# --> directly enable bidirectional attention as all submodules will receive `is_causal = False`.

causal_model = AutoModelForCausalLM.from_pretrained(model_name, device_map=device, attn_implementation="flash_attention_2", torch_dtype="auto").
# it uses the module `Qwen3ModelForCausalLM`, with default value of `is_causal` as `True` 
# --> directly enable causal attention as all submodules will receive `is_causal = True`.

As a result, the end-to-end model has their own attention behavior determined by the default value of is_causal (you can change it by passing is_causal argument into forward though). I know it is not a good idea, and I hope in the future all module will have their own is_causal argument in their forward.

With the kwargs feature, we can enable the bidirectional attention of FA2 by just passing the is_causal into the FA2 function, without further changing the attention module's is_causal.

[vasqu]

My concern is that we gonna have split support for these features, i.e. only flash attention can support this as a simple kwarg handles causality, while other attention flavors (sdpa, eager, flex) will need proper handling of the masks.

Tl;dr: Flash attention is the "simple" case - only change the kwarg, sdpa/eager/flex are more complicated - mask creation

[valencebond]

I'm encountering the same issue. Would it be sufficient to set is_causal = False for all layers to avoid having to manually pass is_causal=False in every forward call? @lucaswychan

for block in model.model.layers:
            block.self_attn.is_causal = False

[vasqu]

@valencebond yes that would already suffice. Note that this will work for FA only, as the mask creation for eager/sdpa/flex etc will need to modify the mask creation as well - just keep that in mind