fix FlashAttentionKwargs rope

Author: garrett361

src/transformers/models/bamba/modeling_bamba.py

@@ -1117,6 +1117,13 @@ def _init_weights(self, module):
                 module.weight.data[module.padding_idx].zero_()
 
 
+def get_position_ids_from_cu_seq_lens(cu_seq_lens: torch.Tensor) -> torch.Tensor:
+    pos_ids = torch.cat(
+        [torch.arange(s, dtype=torch.int32, device=cu_seq_lens.device) for s in cu_seq_lens.diff(dim=-1)], dim=-1
+    )[None]
+    return pos_ids
+
+
 BAMBA_INPUTS_DOCSTRING = r"""
     Args:
         input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
@@ -1281,7 +1288,10 @@ def forward(
         if cache_position is None:
             cache_position = torch.arange(hidden_states.shape[1], device=hidden_states.device)
         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/bamba/modular_bamba.py

@@ -233,6 +233,13 @@ def get_seq_idx_from_cu_seq_lens(cu_seq_lens: torch.Tensor) -> torch.Tensor:
     return seq_idx[None]
 
 
+def get_position_ids_from_cu_seq_lens(cu_seq_lens: torch.Tensor) -> torch.Tensor:
+    pos_ids = torch.cat(
+        [torch.arange(s, dtype=torch.int32, device=cu_seq_lens.device) for s in cu_seq_lens.diff(dim=-1)], dim=-1
+    )[None]
+    return pos_ids
+
+
 # Adapted from transformers.models.mamba2.modeling_mamba2.Mamba2Mixer
 class BambaMixer(nn.Module):
     """
@@ -1029,7 +1036,10 @@ def forward(
         if cache_position is None:
             cache_position = torch.arange(hidden_states.shape[1], device=hidden_states.device)
         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

tests/models/bamba/test_modeling_bamba.py

@@ -22,7 +22,11 @@
 
 from transformers import AutoTokenizer, BambaConfig, is_torch_available
 from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
-from transformers.models.bamba.modular_bamba import get_cu_seq_lens_from_position_ids, get_seq_idx_from_cu_seq_lens
+from transformers.models.bamba.modular_bamba import (
+    get_cu_seq_lens_from_position_ids,
+    get_position_ids_from_cu_seq_lens,
+    get_seq_idx_from_cu_seq_lens,
+)
 from transformers.testing_utils import (
     require_flash_attn,
     require_torch,
@@ -565,9 +569,7 @@ def test_attn_mask_position_ids_flash_attn_equality(self):
             )[None]
 
             torch.testing.assert_close(position_ids_logits, attn_mask_logits_reshaped)
-            # A higher tolerance is needed for the position_ids and FlashAttentionKwargs logits to
-            # match, for unknown reasons.
-            torch.testing.assert_close(position_ids_logits, flash_attn_kwargs_logits, atol=1e-3, rtol=1e-1)
+            torch.testing.assert_close(position_ids_logits, flash_attn_kwargs_logits)
 
 
 @slow
@@ -723,3 +725,24 @@ def test_seq_idx_from_cu_seq_lens() -> None:
     )[None]
     seq_idx_pred = get_seq_idx_from_cu_seq_lens(cu_seq_lens)
     assert torch.allclose(seq_idx_pred, seq_idx)
+
+
+def test_pos_ids_from_cu_seq_lens() -> None:
+    n_chunks = 5
+    max_chunk_len = 64
+
+    seq_lens = torch.randint(1, max_chunk_len, size=(n_chunks,))
+    cu_seq_lens = torch.cat([torch.tensor([0]), seq_lens.cumsum(dim=-1)], dim=-1)
+    pos_ids = torch.cat(
+        [
+            torch.arange(
+                s,
+                dtype=torch.int32,
+                device=cu_seq_lens.device,
+            )
+            for s in cu_seq_lens.diff(dim=-1)
+        ],
+        dim=-1,
+    )[None]
+    pos_ids_pred = get_position_ids_from_cu_seq_lens(cu_seq_lens)
+    assert torch.allclose(pos_ids_pred, pos_ids)