Restore store_kv vs retrieve-only split: use layers[idx] for v5 retrieve path

modeling_sdar.py

@@ -261,21 +261,20 @@ class SDARAttention(nn.Module):
         query_states, key_states = apply_rotary_pos_emb(
             query_states, key_states, cos, sin)
         
-        if past_key_value is not None:
-            # transformers v5 idiom: `.update()` both stores and returns concatenated keys/values.
-            # For v5, we always call `.update()` and ignore the `store_kv` distinction since the cache
-            # handles both paths atomically. For v4 compatibility, we fall back to the old split logic.
+        if past_key_value is not None and kwargs.get("store_kv", False):
+            # Store new kv into the cache (and get concatenated result).
+            key_states, value_states = past_key_value.update(
+                key_states, value_states, self.layer_idx)
+        elif past_key_value is not None and not kwargs.get("store_kv", False) and len(past_key_value) > self.layer_idx:
+            # Retrieve only (don't store): read cached kv and concatenate with current block kv.
+            # transformers v5 DynamicCache uses `.layers[idx].keys/.values`; v4 supports `cache[idx]`.
             if hasattr(past_key_value, "layers"):
-                key_states, value_states = past_key_value.update(
-                    key_states, value_states, self.layer_idx
-                )
-            elif kwargs.get("store_kv", False):
-                key_states, value_states = past_key_value.update(
-                    key_states, value_states, self.layer_idx)
-            elif not kwargs.get("store_kv", False) and len(past_key_value) > self.layer_idx:
+                layer = past_key_value.layers[self.layer_idx]
+                past_key_states, past_value_states = layer.keys, layer.values
+            else:
                 past_key_states, past_value_states = past_key_value[self.layer_idx]
-                key_states = torch.cat([past_key_states, key_states], dim=-2)
-                value_states = torch.cat([past_value_states, value_states], dim=-2)
+            key_states = torch.cat([past_key_states, key_states], dim=-2)
+            value_states = torch.cat([past_value_states, value_states], dim=-2)
         
         attention_mask = attention_mask.bool() if attention_mask is not None else None
         if torch.all(attention_mask):  # decoding