Skip attention_mask.all() GPU-CPU sync during generation#43088
Conversation
During autoregressive generation, `_ignore_causal_mask_sdpa()` calls `padding_mask.all()` on every token to check if the mask contains any padding. This causes a GPU-CPU synchronization that adds ~2-3ms per token, significantly impacting generation latency. This PR adds a context variable `_attention_mask_all_true` that is set at the start of generation based on the initial attention_mask. If the mask is all-True (no padding), the context variable allows subsequent calls to skip the expensive `.all()` check, since during generation we only append ones to the mask. This reduces `.all()` calls from N+1 to 1 for N generated tokens (e.g., from 201 to 1 for 200 tokens). 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <[email protected]>
🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <[email protected]>
| # Check if attention_mask is all-True to avoid per-token GPU-CPU sync in masking utils. | ||
| # During generation, if the mask starts all-True and we only append ones, it stays all-True. | ||
| attention_mask = model_kwargs.get("attention_mask") | ||
| mask_all_true = attention_mask is None or bool(attention_mask.all()) | ||
| mask_token = _attention_mask_all_true.set(mask_all_true) |
There was a problem hiding this comment.
hm, we can also set the mask to None and avoid the sync in masking utils
There was a problem hiding this comment.
meaning, we check once at the start?
if bool(model_kwargs.get("attention_mask").all()):
model_kwargs.set("attention_mask", None)(or otherwise, of course, feel free to implement it as you want, sounds like a very minimal change)
There was a problem hiding this comment.
yeah, once in the beginning. Also want to cc @vasqu for attn masks, maybe we don't even need to tweak generation loop
There was a problem hiding this comment.
Only sdpa and flash attention can skip in those cases. Eager and flex attention (I think) will still need their mask to work correctly - at least when we have a causal model
Note that the way we skip mask creation atm is data-flow dependent, meaning this could indeed help in compiled cases:
SDPA
Flash
transformers/src/transformers/masking_utils.py
Lines 564 to 565 in 136c621
There was a problem hiding this comment.
Thanks @vasqu ! I clearly do not understand the complexities for each attention implementation, but I now added a breakdown per implementation: #43089 (comment)
_update_causal_mask affects eager and sdpa (63 syncs), then _ignore_causal_mask_sdpa affects sdap with 64 syncs, and finally _update_causal_mask affects flash attention with 64 syncs.
My understanding is that if the attention mask is all true, we could ignore it for all these cases
There was a problem hiding this comment.
If everything is true, we can ignore for SDPA and FA - they have a flag to enable proper causality without a mask.
However, flex attention and eager attention cannot do that. They require a mask at all times (at least when we have causal models and these are the majority of models).
So we need to differentiate per attention whether we can avoid these sync
- SDPA + FA: Yes
- Flex + Eager: No
Then we also have a few special models like PaliGemma which have special mixed masks which we would not be able to skip in any case. I assume that should be a small number of models tho
I guess there could be cases where people define a custom mask, but I don't know of any.
Summary
_attention_mask_all_trueto track if attention_mask is known to be all-True during generationgenerate()based on the initial mask state_ignore_causal_mask_sdpa()to skip the expensive.all()check when the context variable indicates no paddingProblem
During autoregressive generation,
_ignore_causal_mask_sdpa()inmasking_utils.pycallspadding_mask.all()on every generated token to check if the mask contains any padding. This causes a GPU-CPU synchronization that adds ~2-3ms per token, significantly impacting generation latency.Solution
Since during generation we only append ones (1s) to the attention mask, if the mask starts all-True (no padding), it will remain all-True throughout generation. We check this once at the start and use a context variable to skip subsequent
.all()checks.Results
.all()calls from N+1 to 1 for N generated tokensTest plan
.all()call count reduced from 201 to 1 for 200 tokens🤖 Generated with Claude Code