You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
Support SDPA & Flash Attention 2 for LayoutLMv3 #35467
Still open: src/transformers/models/layoutlmv3/modeling_layoutlmv3.py doesn't use the attention interface yet (earlier attempts #35469, #41801, #42225 were closed). The path is to port its self-attention to ALL_ATTENTION_FUNCTIONS as BERT does (https://huggingface.co/docs/transformers/main/en/attention_interface); note LayoutLMv3 adds 1D/2D relative position biases to the scores, so SDPA can take them as an additive mask, while flash-attention would need to be disabled when those biases are used.
Plan: move LayoutLMv3SelfAttention onto the attention interface like BERT, and pass the two learned relative position biases as position_bias, the way T5 and UDOP do. The CogView softmax stays in the eager path. SDPA only. Flash attention can't take an additive bias, so I'd skip it and mark test_sdpa_can_dispatch_on_flash as skipped with a reason, like UDOP.
I tried it locally on an RTX 5080. Eager output is bit for bit the same as main at the real head size (64), SDPA matches it to about 1e-6 in fp32, and the LayoutLMv3 tests, including slow ones, pass. With random weights and base-sized shapes SDPA was 1.2x to 1.5x faster for inference and 1.3x to 1.85x for training, with the same or lower peak memory. That is one GPU only.
Two questions before I open a PR:
This makes SDPA the default for LayoutLMv3, so output_attentions=True returns empty attentions unless you load with attn_implementation="eager" (same as BERT). Is that ok here?
PP-DocLayoutV2 inherits these classes through modular, so its generated file changes by a few lines. Fine to include in the same PR?
I used an AI coding assistant for the investigation. I will write the PR myself and run the tests.
Activity
[-]Support Flash Attention 2 for LayoutLMv3[/-][+]Support SDPA & Flash Attention 2 for LayoutLMv3[/+]YashPCHFL commentedon Mar 2, 2025
++ This will drastically improve the speed of LayoutLmv3.
I see the PR is already raised
Add SDPA and FlashAttention-2 support to LayoutLMv3 using unified att…
Add SDPA and FlashAttention-2 support for LayoutLMv3
Add SDPA and FlashAttention-2 support for LayoutLMv3
ArthurZucker commentedon Oct 20, 2026
Still open:
src/transformers/models/layoutlmv3/modeling_layoutlmv3.pydoesn't use the attention interface yet (earlier attempts #35469, #41801, #42225 were closed). The path is to port its self-attention toALL_ATTENTION_FUNCTIONSas BERT does (https://huggingface.co/docs/transformers/main/en/attention_interface); note LayoutLMv3 adds 1D/2D relative position biases to the scores, so SDPA can take them as an additive mask, while flash-attention would need to be disabled when those biases are used.Kash6 commentedon Nov 11, 2026
@ArthurZucker @qgallouedec
Is anyone working on this? I can try something.
Plan: move LayoutLMv3SelfAttention onto the attention interface like BERT, and pass the two learned relative position biases as position_bias, the way T5 and UDOP do. The CogView softmax stays in the eager path. SDPA only. Flash attention can't take an additive bias, so I'd skip it and mark test_sdpa_can_dispatch_on_flash as skipped with a reason, like UDOP.
I tried it locally on an RTX 5080. Eager output is bit for bit the same as main at the real head size (64), SDPA matches it to about 1e-6 in fp32, and the LayoutLMv3 tests, including slow ones, pass. With random weights and base-sized shapes SDPA was 1.2x to 1.5x faster for inference and 1.3x to 1.85x for training, with the same or lower peak memory. That is one GPU only.
Two questions before I open a PR:
This makes SDPA the default for LayoutLMv3, so output_attentions=True returns empty attentions unless you load with attn_implementation="eager" (same as BERT). Is that ok here?
PP-DocLayoutV2 inherits these classes through modular, so its generated file changes by a few lines. Fine to include in the same PR?
I used an AI coding assistant for the investigation. I will write the PR myself and run the tests.