Skip to content

fix FlashAttentionKwargs RoPE #35941

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
garrett361 wants to merge 5 commits into from
Closed

fix FlashAttentionKwargs RoPE #35941

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
08da972
fix FlashAttentionKwargs RoPE
garrett361 Jan 28, 2025
0804ea5
run modular_model_converter.py
garrett361 Jan 28, 2025
06a268e
minimize tensor allocations
garrett361 Feb 5, 2025
973d362
mv get_position_ids_from_cu_seq_lens to flash util
garrett361 Feb 11, 2025
0fb7312
minor import fix
garrett361 Feb 11, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 14 additions & 0 deletions src/transformers/modeling_flash_attention_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -387,3 +387,17 @@ class FlashAttentionKwargs(TypedDict, total=False):
cu_seq_lens_k: Optional[torch.LongTensor]
max_length_q: Optional[int]
max_length_k: Optional[int]


def get_position_ids_from_cu_seq_lens(cu_seq_lens: torch.Tensor) -> torch.Tensor:
if cu_seq_lens.ndim != 1:
raise ValueError(f"cu_seq_lens must be a 1D tensor, received {cu_seq_lens.ndim=}.")
pos_ids = torch.empty(cu_seq_lens[-1], device=cu_seq_lens.device, dtype=torch.int32)
seq_lens = cu_seq_lens.diff(dim=-1)
max_arange = torch.arange(seq_lens.max(), dtype=torch.int32, device=cu_seq_lens.device)
start = torch.tensor(0, device=cu_seq_lens.device, dtype=torch.int32)
for s in seq_lens:
pos_ids[start : start + s] = max_arange[:s]
start += s

return pos_ids[None]
7 changes: 5 additions & 2 deletions src/transformers/models/aria/modeling_aria.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@
from ...cache_utils import Cache, DynamicCache, StaticCache
from ...generation import GenerationMixin
from ...modeling_attn_mask_utils import AttentionMaskConverter
from ...modeling_flash_attention_utils import FlashAttentionKwargs
from ...modeling_flash_attention_utils import FlashAttentionKwargs, get_position_ids_from_cu_seq_lens
from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, ModelOutput
from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS
from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
Expand Down Expand Up @@ -939,7 +939,10 @@ def forward(
)

if position_ids is None:
position_ids = cache_position.unsqueeze(0)
if "cu_seq_lens_q" in flash_attn_kwargs:
position_ids = get_position_ids_from_cu_seq_lens(flash_attn_kwargs["cu_seq_lens_q"])
else:
position_ids = cache_position.unsqueeze(0)

causal_mask = self._update_causal_mask(
attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions
Expand Down
7 changes: 5 additions & 2 deletions src/transformers/models/cohere/modeling_cohere.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@
from ...cache_utils import Cache, DynamicCache, StaticCache
from ...generation import GenerationMixin
from ...modeling_attn_mask_utils import AttentionMaskConverter
from ...modeling_flash_attention_utils import FlashAttentionKwargs
from ...modeling_flash_attention_utils import FlashAttentionKwargs, get_position_ids_from_cu_seq_lens
from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS
from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
Expand Down Expand Up @@ -589,7 +589,10 @@ def forward(
)

if position_ids is None:
position_ids = cache_position.unsqueeze(0)
if "cu_seq_lens_q" in flash_attn_kwargs:
position_ids = get_position_ids_from_cu_seq_lens(flash_attn_kwargs["cu_seq_lens_q"])
else:
position_ids = cache_position.unsqueeze(0)

causal_mask = self._update_causal_mask(
attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions
Expand Down
11 changes: 9 additions & 2 deletions src/transformers/models/diffllama/modeling_diffllama.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,11 @@
from ...cache_utils import Cache, DynamicCache, StaticCache
from ...generation import GenerationMixin
from ...modeling_attn_mask_utils import AttentionMaskConverter
from ...modeling_flash_attention_utils import FlashAttentionKwargs, _flash_attention_forward
from ...modeling_flash_attention_utils import (
FlashAttentionKwargs,
_flash_attention_forward,
get_position_ids_from_cu_seq_lens,
)
from ...modeling_outputs import (
BaseModelOutputWithPast,
CausalLMOutputWithPast,
Expand Down Expand Up @@ -828,7 +832,10 @@ def forward(
)

if position_ids is None:
position_ids = cache_position.unsqueeze(0)
if "cu_seq_lens_q" in flash_attn_kwargs:
position_ids = get_position_ids_from_cu_seq_lens(flash_attn_kwargs["cu_seq_lens_q"])
else:
position_ids = cache_position.unsqueeze(0)

causal_mask = self._update_causal_mask(
attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions
Expand Down
7 changes: 5 additions & 2 deletions src/transformers/models/emu3/modeling_emu3.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@
from ...cache_utils import Cache, DynamicCache, StaticCache
from ...generation import GenerationMixin
from ...modeling_attn_mask_utils import AttentionMaskConverter
from ...modeling_flash_attention_utils import FlashAttentionKwargs
from ...modeling_flash_attention_utils import FlashAttentionKwargs, get_position_ids_from_cu_seq_lens
from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS
from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
Expand Down Expand Up @@ -1407,7 +1407,10 @@ def forward(
)

if position_ids is None:
position_ids = cache_position.unsqueeze(0)
if "cu_seq_lens_q" in flash_attn_kwargs:
position_ids = get_position_ids_from_cu_seq_lens(flash_attn_kwargs["cu_seq_lens_q"])
else:
position_ids = cache_position.unsqueeze(0)

causal_mask = self._update_causal_mask(
attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions
Expand Down
7 changes: 5 additions & 2 deletions src/transformers/models/glm/modeling_glm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@
from ...cache_utils import Cache, DynamicCache, StaticCache
from ...generation import GenerationMixin
from ...modeling_attn_mask_utils import AttentionMaskConverter
from ...modeling_flash_attention_utils import FlashAttentionKwargs
from ...modeling_flash_attention_utils import FlashAttentionKwargs, get_position_ids_from_cu_seq_lens
from ...modeling_outputs import (
BaseModelOutputWithPast,
CausalLMOutputWithPast,
Expand Down Expand Up @@ -570,7 +570,10 @@ def forward(
)

if position_ids is None:
position_ids = cache_position.unsqueeze(0)
if "cu_seq_lens_q" in flash_attn_kwargs:
position_ids = get_position_ids_from_cu_seq_lens(flash_attn_kwargs["cu_seq_lens_q"])
else:
position_ids = cache_position.unsqueeze(0)

causal_mask = self._update_causal_mask(
attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions
Expand Down
7 changes: 5 additions & 2 deletions src/transformers/models/helium/modeling_helium.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@
from ...cache_utils import Cache, DynamicCache, StaticCache
from ...generation import GenerationMixin
from ...modeling_attn_mask_utils import AttentionMaskConverter
from ...modeling_flash_attention_utils import FlashAttentionKwargs
from ...modeling_flash_attention_utils import FlashAttentionKwargs, get_position_ids_from_cu_seq_lens
from ...modeling_outputs import (
BaseModelOutputWithPast,
CausalLMOutputWithPast,
Expand Down Expand Up @@ -557,7 +557,10 @@ def forward(
)

if position_ids is None:
position_ids = cache_position.unsqueeze(0)
if "cu_seq_lens_q" in flash_attn_kwargs:
position_ids = get_position_ids_from_cu_seq_lens(flash_attn_kwargs["cu_seq_lens_q"])
else:
position_ids = cache_position.unsqueeze(0)

causal_mask = self._update_causal_mask(
attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions
Expand Down
7 changes: 5 additions & 2 deletions src/transformers/models/llama/modeling_llama.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@
from ...cache_utils import Cache, DynamicCache, StaticCache
from ...generation import GenerationMixin
from ...modeling_attn_mask_utils import AttentionMaskConverter
from ...modeling_flash_attention_utils import FlashAttentionKwargs
from ...modeling_flash_attention_utils import FlashAttentionKwargs, get_position_ids_from_cu_seq_lens
from ...modeling_outputs import (
BaseModelOutputWithPast,
CausalLMOutputWithPast,
Expand Down Expand Up @@ -559,7 +559,10 @@ def forward(
)

if position_ids is None:
position_ids = cache_position.unsqueeze(0)
if "cu_seq_lens_q" in flash_attn_kwargs:
position_ids = get_position_ids_from_cu_seq_lens(flash_attn_kwargs["cu_seq_lens_q"])
else:
position_ids = cache_position.unsqueeze(0)

causal_mask = self._update_causal_mask(
attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions
Expand Down
7 changes: 5 additions & 2 deletions src/transformers/models/mistral/modeling_mistral.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
from ...cache_utils import Cache, DynamicCache, SlidingWindowCache, StaticCache
from ...generation import GenerationMixin
from ...modeling_attn_mask_utils import AttentionMaskConverter
from ...modeling_flash_attention_utils import FlashAttentionKwargs
from ...modeling_flash_attention_utils import FlashAttentionKwargs, get_position_ids_from_cu_seq_lens
from ...modeling_outputs import (
BaseModelOutputWithPast,
CausalLMOutputWithPast,
Expand Down Expand Up @@ -531,7 +531,10 @@ def forward(
)

if position_ids is None:
position_ids = cache_position.unsqueeze(0)
if "cu_seq_lens_q" in flash_attn_kwargs:
position_ids = get_position_ids_from_cu_seq_lens(flash_attn_kwargs["cu_seq_lens_q"])
else:
position_ids = cache_position.unsqueeze(0)

causal_mask = self._update_causal_mask(
attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions
Expand Down
7 changes: 5 additions & 2 deletions src/transformers/models/olmo/modeling_olmo.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@
from ...cache_utils import Cache, DynamicCache, StaticCache
from ...generation import GenerationMixin
from ...modeling_attn_mask_utils import AttentionMaskConverter
from ...modeling_flash_attention_utils import FlashAttentionKwargs
from ...modeling_flash_attention_utils import FlashAttentionKwargs, get_position_ids_from_cu_seq_lens
from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS
from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
Expand Down Expand Up @@ -535,7 +535,10 @@ def forward(
)

if position_ids is None:
position_ids = cache_position.unsqueeze(0)
if "cu_seq_lens_q" in flash_attn_kwargs:
position_ids = get_position_ids_from_cu_seq_lens(flash_attn_kwargs["cu_seq_lens_q"])
else:
position_ids = cache_position.unsqueeze(0)

causal_mask = self._update_causal_mask(
attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions
Expand Down
7 changes: 5 additions & 2 deletions src/transformers/models/olmo2/modeling_olmo2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
from ...cache_utils import Cache, DynamicCache, StaticCache
from ...generation import GenerationMixin
from ...modeling_attn_mask_utils import AttentionMaskConverter
from ...modeling_flash_attention_utils import FlashAttentionKwargs
from ...modeling_flash_attention_utils import FlashAttentionKwargs, get_position_ids_from_cu_seq_lens
from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS
from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
Expand Down Expand Up @@ -536,7 +536,10 @@ def forward(
)

if position_ids is None:
position_ids = cache_position.unsqueeze(0)
if "cu_seq_lens_q" in flash_attn_kwargs:
position_ids = get_position_ids_from_cu_seq_lens(flash_attn_kwargs["cu_seq_lens_q"])
else:
position_ids = cache_position.unsqueeze(0)

causal_mask = self._update_causal_mask(
attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions
Expand Down
7 changes: 5 additions & 2 deletions src/transformers/models/phi3/modeling_phi3.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@
from ...cache_utils import Cache, DynamicCache, SlidingWindowCache, StaticCache
from ...generation import GenerationMixin
from ...modeling_attn_mask_utils import AttentionMaskConverter
from ...modeling_flash_attention_utils import FlashAttentionKwargs
from ...modeling_flash_attention_utils import FlashAttentionKwargs, get_position_ids_from_cu_seq_lens
from ...modeling_outputs import (
BaseModelOutputWithPast,
CausalLMOutputWithPast,
Expand Down Expand Up @@ -601,7 +601,10 @@ def forward(
)

if position_ids is None:
position_ids = cache_position.unsqueeze(0)
if "cu_seq_lens_q" in flash_attn_kwargs:
position_ids = get_position_ids_from_cu_seq_lens(flash_attn_kwargs["cu_seq_lens_q"])
else:
position_ids = cache_position.unsqueeze(0)

causal_mask = self._update_causal_mask(
attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions
Expand Down
7 changes: 5 additions & 2 deletions src/transformers/models/qwen2/modeling_qwen2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
from ...cache_utils import Cache, DynamicCache, StaticCache
from ...generation import GenerationMixin
from ...modeling_attn_mask_utils import AttentionMaskConverter
from ...modeling_flash_attention_utils import FlashAttentionKwargs
from ...modeling_flash_attention_utils import FlashAttentionKwargs, get_position_ids_from_cu_seq_lens
from ...modeling_outputs import (
BaseModelOutputWithPast,
CausalLMOutputWithPast,
Expand Down Expand Up @@ -544,7 +544,10 @@ def forward(
)

if position_ids is None:
position_ids = cache_position.unsqueeze(0)
if "cu_seq_lens_q" in flash_attn_kwargs:
position_ids = get_position_ids_from_cu_seq_lens(flash_attn_kwargs["cu_seq_lens_q"])
else:
position_ids = cache_position.unsqueeze(0)

causal_mask = self._update_causal_mask(
attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions
Expand Down
Loading