v5-align: remove store_kv; always update() on cached forward

modeling_sdar.py

@@ -261,37 +261,12 @@ class SDARAttention(nn.Module):
         query_states, key_states = apply_rotary_pos_emb(
             query_states, key_states, cos, sin)
         
-        # I-DLM / strict-causal mode (`config.use_regular_causal=True`) uses the standard HF cache
-        # convention: every cached forward commits its KVs. In SDAR block-diffusion mode, callers pass
-        # `store_kv=True` for commit forwards and `store_kv=False` for retrieve-only denoising passes.
-        # Upstream forwards thread `store_kv: Optional[bool] = None` through, so a missing caller kwarg
-        # still arrives here as `None` — treat that as "use config default" rather than False.
-        store_kv = kwargs.get("store_kv", None)
-        if store_kv is None:
-            store_kv = bool(getattr(self.config, "use_regular_causal", False))
-        if past_key_value is not None and store_kv:
-            # Store new kv into the cache (and get concatenated result).
+        # Standard transformers v5 cache convention: when a cache is provided, always `.update()` it.
+        # Callers that want a read-only forward should pass `past_key_values=None`, or use
+        # `DynamicCache.crop(prev_seq_len)` to roll back the append after reading the logits.
+        if past_key_value is not None:
             key_states, value_states = past_key_value.update(
                 key_states, value_states, self.layer_idx)
-        elif past_key_value is not None and not store_kv 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]`.
-            past_key_states = past_value_states = None
-            if hasattr(past_key_value, "layers"):
-                layer = past_key_value.layers[self.layer_idx]
-                past_key_states, past_value_states = layer.keys, layer.values
-            else:
-                try:
-                    past_key_states, past_value_states = past_key_value[self.layer_idx]
-                except (IndexError, TypeError, KeyError):
-                    past_key_states = past_value_states = None
-            # Guard: v5 DynamicCache has a layer slot from construction time but its `.keys` is None
-            # until the first `.update()` call. On the initial cached forward (e.g. plain `generate()`
-            # that hasn't stored this layer's KV yet), there's nothing to concatenate — just use the
-            # current block's KVs as if no cache.
-            if past_key_states is not None and past_value_states is not None:
-                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
 
@@ -373,7 +348,6 @@ class SDARDecoderLayer(GradientCheckpointingLayer):
         past_key_value: Optional[Cache] = None,
         output_attentions: Optional[bool] = False,
         use_cache: Optional[bool] = False,
-        store_kv: Optional[bool] = False,
         cache_position: Optional[torch.LongTensor] = None,
         # necessary, but kept here for BC
         position_embeddings: Optional[Tuple[torch.Tensor,
@@ -391,7 +365,6 @@ class SDARDecoderLayer(GradientCheckpointingLayer):
             past_key_value=past_key_value,
             output_attentions=output_attentions,
             use_cache=use_cache,
-            store_kv=store_kv,
             cache_position=cache_position,
             position_embeddings=position_embeddings,
             **kwargs,
@@ -543,7 +516,6 @@ class SDARModel(SDARPreTrainedModel):
         past_key_values: Optional[Cache] = None,
         inputs_embeds: Optional[torch.FloatTensor] = None,
         use_cache: Optional[bool] = None,
-        store_kv: Optional[bool] = None,
         output_attentions: Optional[bool] = None,
         output_hidden_states: Optional[bool] = None,
         cache_position: Optional[torch.LongTensor] = None,
@@ -610,7 +582,6 @@ class SDARModel(SDARPreTrainedModel):
                 past_key_value=past_key_values,
                 output_attentions=output_attentions,
                 use_cache=use_cache,
-                store_kv=store_kv,
                 cache_position=cache_position,
                 position_embeddings=position_embeddings,
                 **flash_attn_kwargs,