scaled_dot_product_attention throws an error when both attn_mask and is_causal are set

[moskomule]

Hi, thanks for developing a very wonderful project.

I found torch.nn.functional.scaled_dot_product_attention throws an error when both attn_mask and is_causal are set .
But, currently the language_model.py code uses both.

nanoVLM/models/language_model.py

Line 141 in 6ba9082

is_causal=True # LM attention is causal (masked)

A simple fix is to create a causal mask by yourself, but if there's other ways, I want to know.

[andimarafioti]

When does that happen in our codebase?

[moskomule]

When self.spda is True

nanoVLM/models/language_model.py

Lines 136 to 141 in 6ba9082

	if self.sdpa:
	y = torch.nn.functional.scaled_dot_product_attention(
	q, k, v,
	attn_mask=attention_mask,
	dropout_p=self.dropout if self.training else 0.0,
	is_causal=True # LM attention is causal (masked)

[andimarafioti]

We just changed this today with the kv-cache merge, could you try it again? It might be fixed. Still, I never had this issue, so I find it a bit weird since I was running with sdpa

[moskomule]

I’m wondering why it works in your environment, as throwing an error is an expected behavior according to the PyTorch doc (both the latest stable and master).

[AdonaiVera]

Hi @moskomule ,
I didn’t get that error on my side, but I found the issue when I was adding MPS (Apple Silicon) support. It works fine in the latest PyTorch version, but not on Mac yet. When I checked the problem, I found this issue and saw that it triggered several new PRs.

I believe it’s already fixed in the latest PyTorch version, that’s why it works for @andimarafioti and also for me. Could you try updating your PyTorch version? Maybe also check your Python version, it could be related too.

[moskomule]

I use PyTorch 2.7.0 with Python 3.12.

A very simple example can produce the error:

import torch
from torch.nn import functional as F

x = torch.randn(4, 32, 16, device="cuda")
attn_mask = torch.ones(32, 32, device="cuda")
attn_mask = attn_mask.tril()[None]
out = F.scaled_dot_product_attention(x, x, x, attn_mask=attn_mask, is_causal=True)
>  RuntimeError: _scaled_dot_product_attention: Explicit attn_mask should not be set when is_causal=True

[sbrzz]

Hi! I used two different machines (both Windows) and had the problem in just one.

@andimarafioti I didn't found the fix you mentioned, can you share it?

[sbrzz]

@AdonaiVera I checked the latest PyTorch version — attn_mask and is_causal cannot be used together because masking is handled automatically by scaled_dot_product_attention.

Reference:

https://github.com/pytorch/pytorch/blob/e4de93f6a3e342bab34d3757cf90ec0ccc87e168/torch/nn/functional.py#L5823

I haven’t found the solution mentioned by @andimarafioti. Have you?

If no solution turns up, I’d suggest a simple approach such as:

y = torch.nn.functional.scaled_dot_product_attention(
                q, k_exp, v_exp,
                attn_mask=None if is_causal else additive_attn_mask, 
                dropout_p=self.dropout if self.training else 0.0,
                is_causal=is_causal
            )