fix FlashAttentionKwargs RoPE

src/transformers/models/aria/modeling_aria.py

@@ -861,6 +861,20 @@ def forward(self, x, position_ids):
 """
 
 
+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]
+
+
 @add_start_docstrings(
     "The bare AriaText Model outputting raw hidden-states without any specific head on top.",
     ARIA_TEXT_START_DOCSTRING,
@@ -939,7 +953,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

src/transformers/models/cohere/modeling_cohere.py

@@ -511,6 +511,20 @@ def _init_weights(self, module):
 """
 
 
+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]
+
+
 @add_start_docstrings(
     "The bare Cohere Model outputting raw hidden-states without any specific head on top.",
     COHERE_START_DOCSTRING,
@@ -589,7 +603,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

src/transformers/models/diffllama/modeling_diffllama.py

@@ -750,6 +750,20 @@ def forward(self, x, position_ids):
 """
 
 
+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]
+
+
 @add_start_docstrings(
     "The bare DiffLlama Model outputting raw hidden-states without any specific head on top.",
     DIFFLLAMA_START_DOCSTRING,
@@ -828,7 +842,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

src/transformers/models/emu3/modeling_emu3.py

@@ -1258,6 +1258,20 @@ def forward(self, x, position_ids):
         return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
 
 
+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]
+
+
 EMU3_TEXT_INPUTS_DOCSTRING = r"""
     Args:
         input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
@@ -1407,7 +1421,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

src/transformers/models/glm/modeling_glm.py

@@ -492,6 +492,20 @@ def _init_weights(self, module):
 """
 
 
+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]
+
+
 @add_start_docstrings(
     "The bare Glm Model outputting raw hidden-states without any specific head on top.",
     GLM_START_DOCSTRING,
@@ -570,7 +584,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

src/transformers/models/helium/modeling_helium.py

@@ -479,6 +479,20 @@ def _init_weights(self, module):
 """
 
 
+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]
+
+
 @add_start_docstrings(
     "The bare Helium Model outputting raw hidden-states without any specific head on top.",
     HELIUM_START_DOCSTRING,
@@ -557,7 +571,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

src/transformers/models/llama/modeling_llama.py

@@ -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
@@ -1151,6 +1154,20 @@ def forward(
         )
 
 
+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]
+
+
 __all__ = [
     "LlamaForCausalLM",
     "LlamaModel",

src/transformers/models/mistral/modeling_mistral.py

@@ -453,6 +453,20 @@ def _init_weights(self, module):
 """
 
 
+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]
+
+
 @add_start_docstrings(
     "The bare Mistral Model outputting raw hidden-states without any specific head on top.",
     MISTRAL_START_DOCSTRING,
@@ -531,7 +545,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

src/transformers/models/olmo/modeling_olmo.py

@@ -457,6 +457,20 @@ def _init_weights(self, module):
 """
 
 
+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]
+
+
 @add_start_docstrings(
     "The bare Olmo Model outputting raw hidden-states without any specific head on top.",
     OLMO_START_DOCSTRING,
@@ -535,7 +549,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

src/transformers/models/olmo2/modeling_olmo2.py

@@ -458,6 +458,20 @@ def _init_weights(self, module):
 """
 
 
+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]
+
+
 @add_start_docstrings(
     "The bare Olmo2 Model outputting raw hidden-states without any specific head on top.",
     OLMO2_START_DOCSTRING,
@@ -536,7 +550,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

src/transformers/models/phi3/modeling_phi3.py

@@ -523,6 +523,20 @@ def _init_weights(self, module):
 """
 
 
+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]
+
+
 @add_start_docstrings(
     "The bare Phi3 Model outputting raw hidden-states without any specific head on top.",
     PHI3_START_DOCSTRING,
@@ -601,7 +615,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

src/transformers/models/qwen2/modeling_qwen2.py

@@ -466,6 +466,20 @@ def _init_weights(self, module):
 """
 
 
+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]
+
+
 @add_start_docstrings(
     "The bare Qwen2 Model outputting raw hidden-states without any specific head on top.",
     QWEN2_START_DOCSTRING,
@@ -544,7 +558,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