Upload RavenForCausalLM

config.json

@@ -11,9 +11,7 @@
   "bias": false,
   "block_class_name": "SandwichBlock",
   "block_size": 4096,
-  "bos_token_id": 65504,
   "effective_expected_depth": 132,
-  "eos_token_id": 65505,
   "head_dim": 96,
   "init_orthogonal": false,
   "init_strategy": "takase",
@@ -39,7 +37,6 @@
   "norm_class_name": "RMSNorm_llama",
   "norm_eps": 1e-06,
   "num_key_value_heads": 55,
-  "pad_token_id": 65509,
   "padded_vocab_size": 65536,
   "padding_multiple": 4096,
   "qk_bias": true,

generation_config.json

@@ -1,8 +1,4 @@
 {
   "_from_model_config": true,
-  "bos_token_id": 65504,
-  "eos_token_id": 65505,
-  "pad_token_id": 65509,
-  "use_cache": true, 
   "transformers_version": "4.44.2"
 }

raven_modeling_minimal.py

@@ -1,11 +1,11 @@
-"""Minimal modeling.py file for HF compatibility and funny zero-shot experiments. Usability for finetuning not guaranteed"""
+"""Minimal modeling.py file for HF compatibility and funny zero-shot experiments. Use only for inference."""
 
 import torch
 import math
 
 from torch import Tensor
 from dataclasses import dataclass
-from typing import Optional, Union
+from typing import Optional, Union, Any
 
 from .raven_config_minimal import RavenConfig
 from transformers.cache_utils import Cache, DynamicCache
@@ -13,6 +13,10 @@ from transformers.cache_utils import Cache, DynamicCache
 ###################### Huggingface Glue code I ##################################################################
 from transformers import PreTrainedModel
 from transformers.utils import ModelOutput
+from transformers.generation.utils import GenerateDecoderOnlyOutput
+
+import torch.nn.functional as F
+from transformers import GenerationConfig
 
 
 class RavenPreTrainedModel(PreTrainedModel):
@@ -39,7 +43,7 @@ class CausalLMOutputRecurrentLatents(ModelOutput):
     past_key_values: Optional[Cache] = None
     latent_states: Optional[torch.Tensor] = None
     hidden_states: Optional[torch.Tensor] = None
-    attention_maps: Optional[tuple[torch.Tensor, ...]] = None
+    attention_maps: Optional[dict[int, torch.Tensor]] = None
     stats: Optional[dict] = None
 
 
@@ -66,7 +70,7 @@ class RMSNorm(torch.nn.Module):
 
 
 class HuginnDynamicCache(DynamicCache):
-    def __init__(self) -> None:
+    def __init__(self, lookup_strategy: str = "full") -> None:
         super().__init__()
         self._seen_tokens = 0
         self.key_cache: dict[int, dict[int, torch.Tensor]] = {}
@@ -75,14 +79,24 @@ class HuginnDynamicCache(DynamicCache):
         # the cache is held uncoalesced because certain recurrent steps may be missing for some sequence ids if using
         # per-token adaptive compute. In those cases, the "lookup_strategy" determines how to proceed
         # Also, It is critical that the head indices do not overlap with the recurrent iteration indices
+        self.lookup_strategy = lookup_strategy
 
     def update(
         self,
         key_states: torch.Tensor,
         value_states: torch.Tensor,
         step_idx: int,
-        lookup_strategy: str = "latest",
+        lookup_strategy: Optional[str] = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
+        lookup_strategy = self.lookup_strategy if lookup_strategy is None else lookup_strategy
+        if "compress-" in self.lookup_strategy and step_idx > 1:  # hardcode for current model!
+            compression_stage = int(self.lookup_strategy.split("compress-")[1][1:])
+            if "compress-s" in self.lookup_strategy:
+                new_step_idx = (step_idx - 2) % compression_stage + 2
+            else:
+                new_step_idx = (step_idx - 2) // compression_stage + 2
+            # @ print(step_idx, new_step_idx, compression_stage)
+            step_idx = new_step_idx
         # Init
         if step_idx not in self.key_cache:
             self.key_cache[step_idx] = {}
@@ -92,31 +106,49 @@ class HuginnDynamicCache(DynamicCache):
             self._seen_tokens += key_states.shape[-2]
         # Add entries to cache
         for idx, entry in enumerate(key_states.unbind(dim=-2)):
-            assert self._seen_tokens - key_states.shape[-2] + idx not in self.key_cache[step_idx]
+            if "compress-" not in self.lookup_strategy:
+                assert step_idx < 0 or self._seen_tokens - key_states.shape[-2] + idx not in self.key_cache[step_idx]
+            # print(f"Overwrote cache entry for step_idx {step_idx}") # likely the head
             self.key_cache[step_idx][self._seen_tokens - key_states.shape[-2] + idx] = entry
         for idx, entry in enumerate(value_states.unbind(dim=-2)):
             self.value_cache[step_idx][self._seen_tokens - value_states.shape[-2] + idx] = entry
 
         # Materialize past state based on lookup strategy:
-        if len(self.key_cache[step_idx]) == self._seen_tokens:
+        if len(self.key_cache[step_idx]) == self._seen_tokens or self.lookup_strategy == "full":
             # All entries are present, materialize cache as normal
             return (
                 torch.stack(list(self.key_cache[step_idx].values()), dim=-2),
                 torch.stack(list(self.value_cache[step_idx].values()), dim=-2),
             )
         else:  # some entries where not previously computed
-            if lookup_strategy == "latest":
+            # if lookup_strategy.startswith("latest"):
+            #     latest_keys = []
+            #     latest_values = []
+            #     for token_pos in range(self._seen_tokens):
+            #         # Find the latest step that has this token position
+            #         max_step = max((s for s in range(step_idx + 1) if token_pos in self.key_cache[s]), default=None)
+            #         if max_step is None:
+            #             raise ValueError(f"No cache entry found for token position {token_pos}")
+            #         latest_keys.append(self.key_cache[max_step][token_pos])
+            #         latest_values.append(self.value_cache[max_step][token_pos])
+            #     return torch.stack(latest_keys, dim=-2), torch.stack(latest_values, dim=-2)
+            if lookup_strategy.startswith("latest-m4"):
                 latest_keys = []
                 latest_values = []
                 for token_pos in range(self._seen_tokens):
-                    # Find the latest step that has this token position
-                    max_step = max((s for s in range(step_idx + 1) if token_pos in self.key_cache[s]), default=None)
+                    # For steps >= 2, use modulo 4
+                    if step_idx >= 2:
+                        # Find valid steps for this token position
+                        valid_steps = [s for s in range(step_idx + 1) if token_pos in self.key_cache[s]]
+                        max_step = max([s for s in valid_steps if s >= 2 and s % 4 == step_idx % 4])
+                    else:
+                        max_step = step_idx if token_pos in self.key_cache[step_idx] else 0
                     if max_step is None:
                         raise ValueError(f"No cache entry found for token position {token_pos}")
                     latest_keys.append(self.key_cache[max_step][token_pos])
                     latest_values.append(self.value_cache[max_step][token_pos])
                 return torch.stack(latest_keys, dim=-2), torch.stack(latest_values, dim=-2)
-            elif lookup_strategy == "skip":
+            elif lookup_strategy.startswith("skip"):
                 existing_keys = []
                 existing_values = []
                 for token_pos in range(self._seen_tokens):
@@ -124,15 +156,22 @@ class HuginnDynamicCache(DynamicCache):
                         existing_keys.append(self.key_cache[step_idx][token_pos])
                         existing_values.append(self.value_cache[step_idx][token_pos])
                 return torch.stack(existing_keys, dim=-2), torch.stack(existing_values, dim=-2)
-            elif lookup_strategy == "randomized":  # sanity check
+            elif lookup_strategy.startswith("randomized"):  # sanity check
                 rand_keys = []
                 rand_values = []
                 for token_pos in range(self._seen_tokens):
-                    # Find steps that have this token position
-                    steps = [s for s in range(step_idx + 1) if token_pos in self.key_cache[s]]
-                    rand_step = steps[torch.randint(len(steps), (1,))]
-                    rand_keys.append(self.key_cache[rand_step][token_pos])
-                    rand_values.append(self.value_cache[rand_step][token_pos])
+                    if step_idx < 2:  # For prelude steps
+                        max_step = step_idx if token_pos in self.key_cache[step_idx] else 0
+                    else:  # Get all steps from same block position
+                        curr_modulo = (step_idx - 2) % 4 + 2
+                        valid_steps = [
+                            s
+                            for s in range(2, step_idx + 1)
+                            if (s - 2) % 4 + 2 == curr_modulo and token_pos in self.key_cache[s]
+                        ]
+                        max_step = valid_steps[torch.randint(len(valid_steps), (1,))]
+                    rand_keys.append(self.key_cache[max_step][token_pos])
+                    rand_values.append(self.value_cache[max_step][token_pos])
                 return torch.stack(rand_keys, dim=-2), torch.stack(rand_values, dim=-2)
             else:
                 raise ValueError(f"Unknown lookup strategy: {lookup_strategy}")
@@ -146,6 +185,18 @@ class HuginnDynamicCache(DynamicCache):
     def get_seq_length(self, step_idx: int = 0) -> int:
         return self._seen_tokens
 
+    def get_memory_usage(self) -> float:
+        total_bytes = 0
+        # For each recurrent step/layer index
+        for step_idx in self.key_cache:
+            # Get the sequence cache for this step
+            key_seq_cache = self.key_cache[step_idx]
+            for seq_idx in key_seq_cache:
+                key_tensor = key_seq_cache[seq_idx]
+                # Add memory for of key tensors, assuming value is the same
+                total_bytes += key_tensor.nelement() * key_tensor.element_size()
+        return total_bytes * 2 / (1024 * 1024)
+
 
 class CausalSelfAttention(torch.nn.Module):
     def __init__(self, config: RavenConfig) -> None:
@@ -169,7 +220,8 @@ class CausalSelfAttention(torch.nn.Module):
         step_idx: int,
         mask: Optional[Tensor] = None,
         past_key_values: Optional[Cache] = None,
-    ) -> Tensor:
+        return_attn: bool = False,
+    ) -> tuple[Tensor, Optional[Tensor]]:
         B, S, E = x.shape  # batch size, sequence length, embedding dimensionality (n_embd)
         q, k, v = self.Wqkv(x).split(self.chunks, dim=2)
         q = q.view(B, S, self.n_head, self.head_dim)
@@ -189,11 +241,28 @@ class CausalSelfAttention(torch.nn.Module):
         if past_key_values is not None:
             k, v = past_key_values.update(k, v, step_idx)
 
-        y = torch.nn.functional.scaled_dot_product_attention(
-            q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=q.shape[2] > 1
-        )
+        if return_attn:
+            y, attention_map = self.compute_eager_sdpa(q, k, v, attn_mask=mask)
+        else:
+            y = torch.nn.functional.scaled_dot_product_attention(
+                q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=q.shape[2] > 1
+            )
         y = y.transpose(1, 2).reshape(B, S, E).contiguous()  # reshape is a view if possible (it mostly is)
-        return self.proj(y)
+        return self.proj(y), attention_map if return_attn else None
+
+    def compute_eager_sdpa(self, q, k, v, attn_mask):
+        scale = 1.0 / math.sqrt(self.head_dim)
+        scores = torch.matmul(q, k.transpose(-2, -1)) * scale
+
+        if attn_mask is not None:
+            scores = scores + attn_mask
+        if q.shape[2] > 1:
+            causal_mask = torch.triu(torch.ones(q.shape[2], q.shape[2]), diagonal=1).bool()
+            scores.masked_fill_(causal_mask.to(scores.device), float("-inf"))
+
+        attention_weights = torch.nn.functional.softmax(scores, dim=-1)
+        y = torch.matmul(attention_weights, v)
+        return y, attention_weights.max(dim=1)[0]
 
 
 class GatedMLP(torch.nn.Module):
@@ -232,10 +301,12 @@ class SandwichBlock(torch.nn.Module):
         step_idx: int,
         mask: Optional[Tensor] = None,
         past_key_values: Optional[Cache] = None,
-    ) -> Tensor:
-        x = self.norm_2(self.attn(self.norm_1(x), freqs_cis, step_idx, mask, past_key_values) + x)
+        return_attn: bool = False,
+    ) -> tuple[Tensor, Optional[Tensor]]:
+        attn_out, attn_map = self.attn(self.norm_1(x), freqs_cis, step_idx, mask, past_key_values, return_attn)
+        x = self.norm_2(attn_out + x)
         x = self.norm_4(self.mlp(self.norm_3(x)) + x)
-        return x
+        return x, attn_map
 
 
 class RavenForCausalLM(RavenPreTrainedModel):
@@ -318,13 +389,18 @@ class RavenForCausalLM(RavenPreTrainedModel):
 
         if use_cache and past_key_values is None:
             past_key_values = HuginnDynamicCache()
+        attn_maps = {}
+        return_attn = output_details["return_attention"]
 
         # Non-recurrent prelude
         for block_idx, block in enumerate(self.transformer.prelude):
-            input_embeds = block(input_embeds, freqs_cis, block_idx, attention_mask, past_key_values)
+            input_embeds, attn_map = block(
+                input_embeds, freqs_cis, block_idx, attention_mask, past_key_values, return_attn
+            )
+            attn_maps[block_idx] = attn_map
 
         # Main recurrence
-        x, num_steps_no_grad, num_steps_with_grad, xk = self.iterate_forward(
+        x, num_steps_no_grad, num_steps_with_grad, xk, block_idx, attn_maps = self.iterate_forward(
             input_embeds,  # type: ignore
             input_states,
             freqs_cis,
@@ -332,12 +408,14 @@ class RavenForCausalLM(RavenPreTrainedModel):
             attention_mask,
             past_key_values,
             num_steps,
+            attn_maps,
         )
         latent_states = x.clone().detach()
 
         # Coda layers
         for block_idx, block in enumerate(self.transformer.coda, start=1):
-            x = block(x, freqs_cis, -block_idx, attention_mask, past_key_values)
+            x, attn_map = block(x, freqs_cis, -block_idx, attention_mask, past_key_values, return_attn)
+            attn_maps[-block_idx] = attn_map
         x = self.transformer.ln_f(x)
 
         # Prediction head, assuming labels really are labels and not equal to input_ids
@@ -356,7 +434,7 @@ class RavenForCausalLM(RavenPreTrainedModel):
             past_key_values=past_key_values,
             hidden_states=x if output_details["return_head"] else None,
             latent_states=latent_states if output_details["return_latents"] else None,
-            attention_maps=ValueError() if output_details["return_attention"] else None,  # type: ignore
+            attention_maps=attn_maps if output_details["return_attention"] else None,  # type: ignore
             stats=self.get_stats(logits, x, latent_states, xk, input_embeds, num_steps_no_grad, num_steps_with_grad)
             if output_details["return_stats"]
             else None,
@@ -372,9 +450,9 @@ class RavenForCausalLM(RavenPreTrainedModel):
         mask,
         past_key_values: Optional[Cache] = None,
         num_steps: Optional[torch.Tensor] = None,
+        attn_maps: dict = {},
     ):
         x = xk = self.initialize_state(input_embeds) if input_states is None else input_states.clone()
-
         if num_steps is None:
             num_steps_no_grad, num_steps_with_grad = self.randomized_iteration_sampler()  # type: ignore
         elif hasattr(num_steps, "__len__") and len(num_steps) > 1:
@@ -389,20 +467,123 @@ class RavenForCausalLM(RavenPreTrainedModel):
             # and all parameters are always used
             for step in range(num_steps_no_grad):
                 xk = x
-                x, block_idx = self.core_block_forward(xk, input_embeds, freqs_cis, mask, past_key_values, block_idx)
+                x, block_idx, attn_maps = self.core_block_forward(
+                    xk, input_embeds, freqs_cis, mask, past_key_values, block_idx, attn_maps
+                )
 
         for step in range(num_steps_with_grad):
             xk = x
-            x, block_idx = self.core_block_forward(xk, input_embeds, freqs_cis, mask, past_key_values, block_idx)
-        return self.transformer.ln_f(x), num_steps_no_grad, num_steps_with_grad, xk.detach()
+            x, block_idx, attn_maps = self.core_block_forward(
+                xk, input_embeds, freqs_cis, mask, past_key_values, block_idx, attn_maps
+            )
+        return self.transformer.ln_f(x), num_steps_no_grad, num_steps_with_grad, xk.detach(), block_idx, attn_maps
 
     def core_block_forward(
-        self, x, input_embeds, freqs_cis, mask, past_key_values, block_idx: Union[torch.Tensor, int]
+        self,
+        x,
+        input_embeds,
+        freqs_cis,
+        mask,
+        past_key_values,
+        block_idx: Union[torch.Tensor, int],
+        attn_maps: dict = {},
     ):
         x = self.transformer.adapter(torch.cat([x, input_embeds], dim=-1))
         for idx, block in enumerate(self.transformer.core_block, start=1):
-            x = block(x, freqs_cis, block_idx + idx, mask, past_key_values)
-        return x, block_idx + idx
+            x, attn_map = block(x, freqs_cis, block_idx + idx, mask, past_key_values, return_attn=len(attn_maps) > 0)
+            attn_maps[block_idx + idx] = attn_map
+        return x, block_idx + idx, attn_maps
+
+    @torch.no_grad()
+    def iterate_one_step(
+        self,
+        input_embeds,
+        input_states,
+        position_ids: Optional[torch.Tensor] = None,
+        cache_position: Optional[torch.Tensor] = None,
+        block_idx: Union[torch.Tensor, int] = 0,
+        attention_mask: Optional[Tensor] = None,
+        past_key_values: Optional[Cache] = None,
+        attn_maps: dict = {},
+    ):
+        if position_ids is None and cache_position is None:
+            freqs_cis = self.freqs_cis[:, : input_embeds.shape[1]]
+        elif position_ids is not None:
+            freqs_cis = self.freqs_cis.index_select(1, position_ids.squeeze())
+        elif cache_position is not None:
+            freqs_cis = self.freqs_cis[:, cache_position]
+        x, block_idx, attn_maps = self.core_block_forward(
+            input_states, input_embeds, freqs_cis, attention_mask, past_key_values, block_idx, attn_maps
+        )
+        return x, block_idx, attn_maps
+
+    def predict_from_latents(
+        self,
+        latents,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        cache_position: Optional[torch.Tensor] = None,
+        past_key_values: Optional[Cache] = None,
+        return_attn: bool = False,
+        attn_maps: dict = {},
+    ):
+        if position_ids is None and cache_position is None:
+            freqs_cis = self.freqs_cis[:, : latents.shape[1]]
+        elif position_ids is not None:
+            freqs_cis = self.freqs_cis.index_select(1, position_ids.squeeze())
+        elif cache_position is not None:
+            freqs_cis = self.freqs_cis[:, cache_position]
+        x = self.transformer.ln_f(latents)
+        # Coda layers
+        for block_idx, block in enumerate(self.transformer.coda, start=1):
+            x, attn_map = block(x, freqs_cis, -block_idx, attention_mask, past_key_values)
+        attn_maps[block_idx] = attn_map
+        x = self.transformer.ln_f(x)
+
+        logits = self.lm_head(x).float()
+
+        return CausalLMOutputRecurrentLatents(
+            loss=torch.as_tensor(0.0),
+            log_ppl=torch.as_tensor(0.0),
+            logits=logits,
+            past_key_values=past_key_values,
+            attention_maps=attn_maps if len(attn_maps) > 0 else None,
+        )
+
+    def embed_inputs(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        past_key_values: Optional[Cache] = None,
+        use_cache: bool = False,
+        cache_position: Optional[torch.Tensor] = None,
+        return_attn: bool = False,
+        **kwargs,
+    ) -> tuple[torch.Tensor, int, dict[int, Tensor]]:
+        # Support multiple position formats:
+        if position_ids is None and cache_position is None:
+            freqs_cis = self.freqs_cis[:, : input_ids.shape[1]]
+        elif position_ids is not None:
+            freqs_cis = self.freqs_cis.index_select(1, position_ids.squeeze())
+        elif cache_position is not None:
+            freqs_cis = self.freqs_cis[:, cache_position]
+
+        input_embeds = self.transformer.wte(input_ids)
+
+        if self.emb_scale != 1:
+            input_embeds = input_embeds * self.emb_scale  # type: ignore
+
+        if use_cache and past_key_values is None:
+            past_key_values = HuginnDynamicCache()
+
+        # Non-recurrent prelude
+        attn_maps = {}
+        for block_idx, block in enumerate(self.transformer.prelude):
+            input_embeds, attn_maps = block(
+                input_embeds, freqs_cis, block_idx, attention_mask, past_key_values, return_attn
+            )
+        return input_embeds, block_idx, attn_maps
 
     @torch._dynamo.disable(recursive=False)  # type: ignore
     def randomized_iteration_sampler(self) -> tuple[torch.Tensor, torch.Tensor]:
@@ -462,6 +643,278 @@ class RavenForCausalLM(RavenPreTrainedModel):
                 model_inputs[key] = value
         return model_inputs
 
+    @torch.no_grad()
+    def generate_minimal(
+        self,
+        input_ids: torch.LongTensor,
+        generation_config: Optional[GenerationConfig] = None,  # type: ignore
+        tokenizer=None,
+        streamer=None,
+        continuous_compute=False,  # warm-start state / continuous CoT
+        cache_kwargs: dict = {},
+        **model_kwargs,
+    ) -> Union[torch.Tensor, dict[str, Any]]:
+        """Minimal single-sequence generation. Template for more complicated generate tasks"""
+        # Setup
+        if generation_config is None:
+            generation_config: GenerationConfig = self.generation_config  # type: ignore
+        model_kwargs["past_key_values"] = HuginnDynamicCache(**cache_kwargs)
+        model_kwargs["use_cache"] = True
+        model_kwargs = self._get_initial_cache_position(input_ids, model_kwargs)
+        stop_tokens = self._get_stops(generation_config, tokenizer).to(input_ids.device)
+        if continuous_compute:
+            embedded_inputs, _, _ = self.embed_inputs(input_ids)
+            model_kwargs["input_states"] = self.initialize_state(embedded_inputs)
+        # Generate tokens
+        for _ in range(generation_config.max_length - input_ids.shape[1]):
+            # Forward pass
+            model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
+            outputs = self(**model_inputs)
+            next_token_logits = outputs.logits[0, -1, :]
+            if continuous_compute:
+                current_last_latent = outputs.latent_states[:, -1:, :]
+
+            # Sample or select next token
+            if generation_config.do_sample:
+                if generation_config.temperature:
+                    next_token_logits = next_token_logits / generation_config.temperature
+
+                probs = F.softmax(next_token_logits, dim=-1)
+
+                # Apply top_k
+                if generation_config.top_k:
+                    top_k_probs, _ = torch.topk(probs, generation_config.top_k)
+                    probs[probs < top_k_probs[-1]] = 0
+                # Apply top_p
+                if generation_config.top_p:
+                    sorted_probs = torch.sort(probs, descending=True)[0]
+                    cumsum = torch.cumsum(sorted_probs, dim=-1)
+                    probs[cumsum > generation_config.top_p] = 0
+                # Apply min_p
+                if generation_config.min_p:
+                    probs[probs < generation_config.min_p * probs.max()] = 0
+
+                probs = probs / probs.sum()
+                next_token = torch.multinomial(probs, num_samples=1)
+            else:
+                next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True)
+
+            input_ids = torch.cat([input_ids, next_token[None, :]], dim=-1)  # type: ignore
+
+            if streamer:
+                streamer.put(next_token.cpu())
+
+            # Update model kwargs
+            model_kwargs = self._update_model_kwargs_for_generation(outputs, model_kwargs)
+            if continuous_compute:
+                model_kwargs["input_states"] = current_last_latent
+
+            # Check if we hit a stop token
+            if stop_tokens is not None and next_token in stop_tokens:
+                break
+
+        if streamer:
+            streamer.end()
+
+        if generation_config.return_dict_in_generate:
+            return GenerateDecoderOnlyOutput(
+                sequences=input_ids,
+                scores=None,
+                logits=None,
+                attentions=None,
+                hidden_states=None,
+                past_key_values=model_kwargs.get("past_key_values"),
+            )
+        return input_ids
+
+    @torch.no_grad()
+    def generate_with_adaptive_compute(
+        self,
+        input_ids: torch.LongTensor,
+        generation_config: Optional[GenerationConfig] = None,  # type: ignore
+        tokenizer=None,
+        streamer=None,
+        continuous_compute=False,  # warm-start state / continuous CoT
+        latent_dampening=False,
+        criterion="entropy-diff",
+        exit_threshold: Union[str, float, int] = "auto",
+        cache_kwargs: dict = {},
+        **model_kwargs,
+    ) -> Union[torch.Tensor, GenerateDecoderOnlyOutput]:
+        """Minimal single-sequence generation. Template for more complicated generate tasks"""
+        # Setup
+        if generation_config is None:
+            generation_config: GenerationConfig = self.generation_config  # type: ignore
+        model_kwargs["past_key_values"] = HuginnDynamicCache(**cache_kwargs)
+        model_kwargs["use_cache"] = True
+        model_kwargs = self._get_initial_cache_position(input_ids, model_kwargs)
+        stop_tokens = self._get_stops(generation_config, tokenizer).to(input_ids.device)
+        if continuous_compute:
+            embedded_inputs, _, _ = self.embed_inputs(input_ids)
+            current_last_latent = self.initialize_state(embedded_inputs)
+        compute_steps = []
+
+        # Generate tokens
+        for step in range(generation_config.max_length - input_ids.shape[1]):
+            # Adaptive compute forward
+            model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
+            aux_inputs = {
+                k: model_inputs[k] for k in ["cache_position", "past_key_values", "attention_mask"] if k in model_inputs
+            }
+            embedded_inputs, block_idx, _ = self.embed_inputs(model_inputs["input_ids"], **aux_inputs)
+            if not continuous_compute:
+                current_latents = self.initialize_state(embedded_inputs, deterministic=False)
+            else:
+                current_latents = current_last_latent
+
+            # Prep criterions:
+            if criterion == "entropy-diff":
+                entropy = torch.tensor(100.0, device=input_ids.device)
+                exit_threshold = 1e-3 if exit_threshold == "auto" else float(exit_threshold)
+            elif criterion in ["latent-diff", "none"]:
+                exit_threshold = 0.03 if exit_threshold == "auto" else float(exit_threshold)
+            elif "kl" in criterion:
+                V = self.config.padded_vocab_size
+                log_probs = (1 / V * torch.ones(V, device=input_ids.device)).log()
+                if criterion == "minp-kl":
+                    exit_threshold = 1e-6 if exit_threshold == "auto" else float(exit_threshold)
+                else:
+                    exit_threshold = 5e-4 if exit_threshold == "auto" else float(exit_threshold)
+            elif criterion == "argmax-stability":
+                stable_for_n_steps = 0
+                current_argmax = torch.tensor(-1, dtype=torch.long, device=input_ids.device)
+                exit_threshold = 5 if exit_threshold == "auto" else int(exit_threshold)
+            else:
+                raise ValueError("Invalid adaptive compute strategy.")
+
+            all_latents = []
+            exit_values = []
+            for compute_step in range(model_inputs["num_steps"]):
+                prev_latents = current_latents.clone()
+                current_latents, block_idx, _ = self.iterate_one_step(
+                    embedded_inputs, current_latents, block_idx=block_idx, **aux_inputs
+                )
+                all_latents.append(current_latents if latent_dampening else None)
+                if step > 0:  # do not exit in prefill:
+                    if criterion == "entropy-diff":
+                        prev_entropy = entropy.clone()
+                        outputs = self.predict_from_latents(current_latents, **aux_inputs)
+                        probs = F.softmax(outputs.logits[:, -1, :], dim=-1)  # type: ignore
+                        entropy = -torch.sum(probs * torch.log(probs + 1e-10), dim=-1).mean()
+                        entropy_diff = (entropy - prev_entropy).abs()
+                        exit_values.append(entropy_diff.item())
+                        if entropy_diff < exit_threshold:
+                            break
+                    elif criterion == "latent-diff":
+                        norm_diff = (prev_latents - current_latents).norm() / current_latents.norm()
+                        exit_values.append(norm_diff.item())
+                        if norm_diff < exit_threshold:
+                            break
+                    elif criterion == "kl":
+                        prev_log_probs = log_probs.clone()
+                        outputs = self.predict_from_latents(current_latents, **aux_inputs)
+                        log_probs = F.log_softmax(outputs.logits[:, -1, :], dim=-1)  # type: ignore
+                        kl = F.kl_div(log_probs, prev_log_probs, reduction="none", log_target=True).sum(dim=-1)
+                        exit_values.append(kl.item())
+                        if kl < exit_threshold:
+                            break
+                    elif criterion == "minp-kl":
+                        prev_log_probs = log_probs.clone()
+                        outputs = self.predict_from_latents(current_latents, **aux_inputs)
+                        probs = F.softmax(outputs.logits[:, -1, :], dim=-1)  # type: ignore
+                        probs[probs < 0.1 * probs.max()] = 1 / V
+                        probs = probs / probs.sum()
+                        log_probs = probs.log()
+                        kl = F.kl_div(log_probs, prev_log_probs, reduction="none", log_target=True).sum(dim=-1)
+                        exit_values.append(kl.item())
+                        if kl < exit_threshold:
+                            break
+                    elif criterion == "argmax-stability":
+                        prev_argmax = current_argmax.clone()
+                        outputs = self.predict_from_latents(current_latents, **aux_inputs)
+                        current_argmax = outputs.logits[0, -1, :].argmax(dim=-1)  # type: ignore
+                        if current_argmax == prev_argmax:
+                            stable_for_n_steps += 1
+                        else:
+                            stable_for_n_steps = 0
+                        exit_values.append(stable_for_n_steps)
+                        if stable_for_n_steps >= exit_threshold:
+                            break
+                    elif criterion == "none":
+                        pass
+
+            else:
+                if not latent_dampening:
+                    outputs = self.predict_from_latents(current_latents, **aux_inputs)
+                else:
+                    dampened_latents = torch.sum(torch.cat(all_latents, dim=0), dim=0, keepdim=True)
+                    outputs = self.predict_from_latents(dampened_latents, **aux_inputs)
+            compute_steps.append([compute_step + 1, exit_values])
+
+            next_token_logits = outputs.logits[0, -1, :]  # type: ignore
+            if continuous_compute:  # Save last latent
+                current_last_latent = current_latents[:, -1:, :]
+
+            # Sample or select next token
+            if generation_config.do_sample:
+                if generation_config.temperature:
+                    next_token_logits = next_token_logits / generation_config.temperature
+
+                probs = F.softmax(next_token_logits, dim=-1)
+                # Apply top_k
+                if generation_config.top_k:
+                    top_k_probs, _ = torch.topk(probs, generation_config.top_k)
+                    probs[probs < top_k_probs[-1]] = 0
+                # Apply top_p
+                if generation_config.top_p:
+                    sorted_probs = torch.sort(probs, descending=True)[0]
+                    cumsum = torch.cumsum(sorted_probs, dim=-1)
+                    probs[cumsum > generation_config.top_p] = 0
+                # Apply min_p
+                if generation_config.min_p:
+                    probs[probs < generation_config.min_p * probs.max()] = 0
+
+                probs = probs / probs.sum()
+                next_token = torch.multinomial(probs, num_samples=1)
+            else:
+                next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True)
+
+            input_ids = torch.cat([input_ids, next_token[None, :]], dim=-1)  # type: ignore
+
+            if streamer:
+                streamer.put(next_token.cpu())
+
+            # Update model kwargs
+            model_kwargs = self._update_model_kwargs_for_generation(outputs, model_kwargs)
+
+            # Check if we hit a stop token
+            if stop_tokens is not None and next_token in stop_tokens:
+                break
+
+        if streamer:
+            streamer.end()
+
+        if generation_config.return_dict_in_generate:
+            return GenerateDecoderOnlyOutput(
+                sequences=input_ids,
+                scores=compute_steps,  # type: ignore
+                logits=None,
+                attentions=None,
+                hidden_states=None,
+                past_key_values=model_kwargs.get("past_key_values"),
+            )
+        return input_ids
+
+    def _get_stops(self, generation_config, tokenizer):
+        stop_tokens = set()
+        if generation_config.eos_token_id is not None:
+            stop_tokens.add(generation_config.eos_token_id)
+        if hasattr(generation_config, "stop_strings") and tokenizer and generation_config.stop_strings:
+            for s in generation_config.stop_strings:
+                token_id = tokenizer(s, add_special_tokens=False)["input_ids"][0]
+                stop_tokens.add(token_id)
+        return torch.tensor(list(stop_tokens))
+
     def get_stats(self, logits, x, latent_states, xk, input_embeds, num_steps_no_grad, num_steps_with_grad):
         probs = torch.softmax(logits.float(), dim=-1)
         prob_entropy = torch.where(probs > 0, -probs * probs.log(), 0).sum(dim=-1)