Enable upstream-compatible recurrent caching with dynamic allocation

README.md

@@ -175,7 +175,7 @@ print(tokenizer.decode(out[0]))
 
 `generation_config.json` ships with `eos_token_id = 0` (the EOD byte) and `pad_token_id = 1` so `model.generate()` stops naturally at the trained end-of-document token.
 
-Generation currently uses full recomputation rather than the experimental recurrent cache. `use_cache=True` requests safely fall back to the same numerically verified path, avoiding variant-dependent cache divergence and context-sized preallocation.
+Generation uses recurrent MHA/Hyena state for unpadded prompts. Padded batches automatically use the exact full-recomputation path. Attention KV storage grows with the decoded sequence instead of preallocating the model's full configured context.
 
 ### Attention weights
 
@@ -230,7 +230,7 @@ This HuggingFace port has **not** been tested for fine-tuning - it's verified on
   - `"eager"`: textbook einsum + softmax(QK^T) + einsum. Slowest, used internally when `output_attentions=True` so the attention matrix is materialized.
 - **Block dispatch (`hyena.py`).** StripedHyena 2 has 4 block types, dispatched by `layer_idx` membership in four config lists: `attn_layer_idxs` (MHA + RoPE), `hcl_layer_idxs` (modal-form IIR via FFT), `hcm_layer_idxs` (medium FIR cascade, inner length 128), `hcs_layer_idxs` (short FIR cascade, inner length 7). The disjoint union must equal `range(num_layers)`.
 - **`TELinear` with pure-PyTorch fallback (`layers.py`).** Hyena cascade blocks use a TransformerEngine-backed input projection (3x hidden_size output) that supports FP8 quantisation. When TE is not installed, a `TELinear` fallback class with the same state_dict layout (`weight`, `bias`) is used - checkpoints are cross-loadable.
-- **Generation.** The experimental recurrent cache classes remain available for development, but public generation uses full recomputation because it is numerically consistent across every released variant and avoids context-sized cache allocation.
+- **Generation cache.** Unpadded prompts use recurrent MHA KV plus HCL/HCM/HCS FIR/IIR state matching the upstream Vortex cache path. Padded batches use full recomputation. KV tensors grow geometrically with the decoded length instead of allocating `max_seqlen` up front.
 - **Tokenizer (`tokenization_evo2.py`).** Byte-level UTF-8, vocab_size = 512. Pad token = byte `\x01`. EOS = byte `\x00` (set as `eos_token_id` in `generation_config.json`). Tokenizer does not add EOS at encoding time - matches the original Evo 2 inference pipeline.
 - **Dependencies.** `torch`, `transformers`, `numpy`, `safetensors`, `huggingface_hub`. `transformer-engine[pytorch]` is required for this variant's FP8 input projections. `accelerate` is optional but recommended if you want to load with `device_map="auto"` for multi-GPU sharding. `flash_attn` is optional (only needed if you pass `attn_implementation="flash_attention_2"`).
 

cache.py

@@ -114,6 +114,7 @@ class Evo2Cache:
         )
         self.input_ids: Optional[Tensor] = None
         self.attention_mask: Optional[Tensor] = None
+        self.mode: Optional[str] = None
 
     @property
     def seqlen_offset(self) -> int:
@@ -147,6 +148,7 @@ class Evo2Cache:
         self.hcs.reset()
         self.input_ids = None
         self.attention_mask = None
+        self.mode = None
 
     def prepare_inputs(
         self,
@@ -154,6 +156,7 @@ class Evo2Cache:
         attention_mask: Optional[Tensor],
     ) -> tuple[Tensor, Tensor, bool]:
         """Return full history and whether the caller supplied only a suffix."""
+        self.mode = "recompute"
         requested_mask = (
             attention_mask.ne(0)
             if attention_mask is not None

modeling_evo2.py

@@ -8,9 +8,8 @@ Two top-level model classes:
   * Evo2ForCausalLM  -- with the tied LM head + GenerationMixin so
                         model.generate() works out of the box.
 
-The custom Evo2Cache implementation is retained for future development, but
-generation currently uses full recomputation because cached recurrence is not
-numerically equivalent for every released variant.
+Evo2Cache uses optimized recurrent MHA/Hyena state for unpadded prompts and
+automatically uses exact full recomputation for padded batches.
 
 Per-block dtype: bfloat16 for everything except the modal-form filter
 parameters (`log_poles`, `residues`) which must stay fp32 for stability.
@@ -136,21 +135,48 @@ class StripedHyena2(nn.Module):
                 past_key_values = self.initialize_inference_params(
                     max_batch_size=input_ids.shape[0],
                 )
-            input_ids, attention_mask, suffix_only = past_key_values.prepare_inputs(
-                input_ids,
-                attention_mask,
+            has_padding = (
+                attention_mask is not None
+                and not attention_mask.bool().all()
             )
+            if past_key_values.mode is None:
+                past_key_values.mode = (
+                    "recompute" if has_padding else "recurrent"
+                )
+            if past_key_values.mode == "recurrent" and has_padding:
+                raise ValueError(
+                    "A recurrent Evo2 cache cannot switch to padded inputs. "
+                    "Reset the cache and prefill the padded batch again."
+                )
+            if past_key_values.mode == "recurrent":
+                cached_length = past_key_values.get_seq_length()
+                if cached_length > 0 and input_ids.shape[1] > cached_length:
+                    input_ids = input_ids[:, cached_length:]
+                    if attention_mask is not None:
+                        attention_mask = attention_mask[:, cached_length:]
+            else:
+                input_ids, attention_mask, suffix_only = past_key_values.prepare_inputs(
+                    input_ids,
+                    attention_mask,
+                )
 
         last_hidden, past_kv, hidden_states, attentions = self._forward_features(
             input_ids,
-            inference_params_dict=None,
+            inference_params_dict=(
+                past_key_values
+                if use_cache and past_key_values.mode == "recurrent"
+                else None
+            ),
             attention_mask=attention_mask,
             output_hidden_states=output_hidden_states,
             output_attentions=output_attentions,
         )
         if use_cache:
-            past_key_values.update_history(input_ids, attention_mask)
             past_kv = past_key_values
+            if past_key_values.mode == "recurrent":
+                past_key_values.advance(input_ids.shape[1])
+            else:
+                past_key_values.update_history(input_ids, attention_mask)
             if suffix_only:
                 last_hidden = last_hidden[:, -requested_length:]
                 if output_hidden_states: