Why `is_causal` is not used in `flash_attention_forward` ?
Author: lucaswychan
Pull request: no
Labels: name: Flash Attention

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` in`modeling_qwen3.py`):

```python
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`:

```python
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!!

Comments:
cc @arthurzucker @Cyrilvallez @vasqu for FA2

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

@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 !!!

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

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.

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:

```python3
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`.

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

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
```

@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