Update raven_modeling_minimal.py

raven_modeling_minimal.py

@@ -6,9 +6,10 @@ import math
 from torch import Tensor
 from torch.nn.attention.flex_attention import create_block_mask, BlockMask, flex_attention
 from torch.nn.attention import bias as attn_bias
+from torch.utils.checkpoint import checkpoint
 from dataclasses import dataclass
-from typing import Union, Optional, Any
-
+from typing import Union, Optional, Any, Tuple, Callable, List
+from functools import cache, cached_property
 
 from .raven_config_minimal import RavenConfig
 from transformers.cache_utils import Cache, DynamicCache, StaticCache
@@ -21,8 +22,6 @@ from transformers.generation.utils import GenerateDecoderOnlyOutput
 import torch.nn.functional as F
 from transformers import GenerationConfig
 
-torch.backends.cuda.enable_math_sdp(False)
-
 
 class RavenPreTrainedModel(PreTrainedModel):
     config_class = RavenConfig
@@ -38,9 +37,77 @@ class RavenPreTrainedModel(PreTrainedModel):
     _supports_static_cache = True
     _tp_plan = {}
 
+    @cache
+    def _init_func(self, dim, num_layers):
+        return {
+            "std": math.sqrt(2 / (5 * dim)),
+            "out_proj": math.sqrt(2 / (5 * dim)) / math.sqrt(2 * num_layers),
+            "embedding": math.sqrt(2 / (5 * dim)),
+            "embed_scale": math.sqrt(dim),
+        }
+
+    @property
+    def emb_scale(self):
+        return self._init_func(self.config.n_embd, self.config.effective_expected_depth)["embed_scale"]
+
+    def _normal_(self, tensor, std):
+        return torch.nn.init.trunc_normal_(tensor, mean=0.0, std=std, a=-3 * std, b=3 * std)
+
+    @torch.no_grad()
+    def init_qkv(self, qkv_tensor, init_fn, qk_std, v_std, dim, head_dim):
+        s = qkv_tensor.shape[0]
+        n_kv_heads = (s - dim) // (2 * head_dim)
+        shapes = [dim, n_kv_heads * head_dim, n_kv_heads * head_dim]
+
+        Q, K, V = (
+            qkv_tensor.new_empty([shapes[0], dim]),
+            qkv_tensor.new_empty([shapes[1], dim]),
+            qkv_tensor.new_empty([shapes[2], dim]),
+        )
+        init_fn(Q, qk_std)
+        init_fn(K, qk_std)
+        init_fn(V, v_std)
+        qkv_tensor.data.copy_(torch.cat([Q, K, V], dim=0).contiguous())
+
+    @torch.no_grad()
+    def init_glu(self, glu_tensor, init_fn, w1_std, w2_std):
+        g, h = glu_tensor.shape
+        W1, W2 = (
+            glu_tensor.new_empty([g // 2, h]),
+            glu_tensor.new_empty([g // 2, h]),
+        )
+        init_fn(W1, w1_std)
+        init_fn(W2, w2_std)
+        glu_tensor.data.copy_(torch.cat([W1, W2], dim=0).contiguous())
+
+    @cached_property
+    def _full_name_of_module_lookup(self):
+        return {id(m): n for n, m in self.named_modules()}
+
+    @torch.no_grad()
     def _init_weights(self, module):
-        if not torch.rand((1,)).is_meta:
-            print("Random Initialization not implemented.")
+        _init_values = self._init_func(self.config.n_embd, self.config.effective_expected_depth)
+        name = self._full_name_of_module_lookup[id(module)]
+        if isinstance(module, RMSNorm):
+            torch.nn.init.ones_(module.weight)
+        elif isinstance(module, torch.nn.Linear):
+            if "Wqkv" in name:
+                self.init_qkv(
+                    module.weight,
+                    self._normal_,
+                    float(_init_values["std"]),
+                    float(_init_values["std"]),
+                    self.config.n_embd,
+                    self.config.head_dim,
+                )
+            elif "fc" in name:
+                self.init_glu(module.weight, self._normal_, float(_init_values["std"]), float(_init_values["out_proj"]))
+            elif "mlp.proj" in name or "attn.proj" in name:
+                self._normal_(module.weight, std=float(_init_values["out_proj"]))
+            elif "adapter" in name or "lm_head" in name:
+                self._normal_(module.weight, std=float(_init_values["std"]))
+        elif isinstance(module, torch.nn.Embedding):
+            self._normal_(module.weight, std=float(_init_values["embedding"]))
 
 
 @dataclass
@@ -468,6 +535,9 @@ class SandwichBlock(torch.nn.Module):
         return x
 
 
+#################################### Main Model ##################################################################
+
+
 class RavenForCausalLM(RavenPreTrainedModel, GenerationMixin):
     freqs_cis: torch.Tensor
 
@@ -498,13 +568,15 @@ class RavenForCausalLM(RavenPreTrainedModel, GenerationMixin):
                 ln_f=RMSNorm(config.n_embd, eps=config.norm_eps),  # used twice :>
             )
         )
-        self.emb_scale = config.init_values["embed_scale"]
         # Head
         self.lm_head = torch.nn.Linear(config.n_embd, config.padded_vocab_size, bias=False)
         if self.config.tie_embeddings:
             self.tie_weights()
         # rope
         self.register_buffer("freqs_cis", self._precompute_freqs_cis(), persistent=True)
+        self.gradient_checkpointing = False
+        # Call weight init through HF post init:
+        self.post_init()
 
     def get_input_embeddings(self):
         return self.transformer.wte
@@ -513,11 +585,9 @@ class RavenForCausalLM(RavenPreTrainedModel, GenerationMixin):
         return self.lm_head
 
     def _precompute_freqs_cis(self):
-        # can actually be a buffer now, and remains in fp32! (at least in the settings I tested)
-        freqs_cis = precompute_freqs_cis(
+        return precompute_freqs_cis(
             self.config.n_embd // self.config.num_attention_heads, self.config.block_size, self.config.rope_base, 1
         )
-        return freqs_cis
 
     def compile_mask(
         self,
@@ -557,72 +627,7 @@ class RavenForCausalLM(RavenPreTrainedModel, GenerationMixin):
             H=None,
             Q_LEN=seq_len,
             KV_LEN=kv_length,
-            device=input_ids.device,
-        )
-
-        # # Define mask_mod function
-        # def mask_mod(b, h, q_idx, kv_idx):
-        #     # Always apply causal constraint
-        #     is_causal = q_idx >= kv_idx
-
-        #     # Handle cache vs current tokens
-        #     is_cache = kv_idx < cache_len
-        #     current_idx = kv_idx - cache_len
-
-        #     # For cache: always valid; For current: check padding
-        #     not_pad = input_ids[b, current_idx] != pad_token_id
-        #     valid = is_cache | not_pad
-
-        #     # Apply attention mask if provided
-        #     if attention_mask is not None:
-        #         q_idx_curr = q_idx - cache_len
-        #         attn_valid = attention_mask[b, q_idx_curr, current_idx]
-        #         valid = valid & (is_cache | attn_valid)
-
-        #     return is_causal & valid
-
-        # def mask_mod(b, h, q_idx, kv_idx):
-        #     is_causal = q_idx >= kv_idx
-        #     is_current = (kv_idx >= cache_len) & (kv_idx < kv_length)
-        #     current_idx = kv_idx - cache_len
-
-        #     is_valid = (~is_current) | (
-        #         (current_idx >= 0) & (current_idx < seq_len) & (input_ids != pad_token_id)[b, current_idx % seq_len]
-        #     )
-
-        #     return is_causal & is_valid
-
-        # # Define mask_mod function
-        # def mask_mod(b, h, q_idx, kv_idx):
-        #     # Always apply causal constraint
-        #     is_causal = q_idx >= kv_idx
-
-        #     # Handle cache vs current tokens
-        #     is_cache = kv_idx < cache_len
-        #     current_idx = kv_idx - cache_len
-        #     in_bounds = (current_idx >= 0) & (current_idx < seq_len)
-
-        #     # For cache: always valid; For current: check padding
-        #     not_pad = (input_ids[b, current_idx % seq_len] != pad_token_id) | ~in_bounds
-        #     valid = is_cache | (not_pad & in_bounds)
-
-        #     # Apply attention mask if provided
-        #     if attention_mask is not None:
-        #         q_idx_curr = q_idx - cache_len
-        #         q_in_bounds = (q_idx_curr >= 0) & (q_idx_curr < seq_len)
-        #         attn_valid = attention_mask[b, q_idx_curr % seq_len, current_idx % seq_len] | ~(in_bounds & q_in_bounds)
-        #         valid = valid & (is_cache | attn_valid)
-
-        #     return is_causal & valid
-
-        # Create block mask
-        block_mask = create_block_mask(
-            mask_mod,
-            B=batch_size,
-            H=None,
-            Q_LEN=seq_len,
-            KV_LEN=kv_length,
-            device=input_ids.device,
+            device=str(input_ids.device),
         )
 
         return block_mask
@@ -748,7 +753,7 @@ class RavenForCausalLM(RavenPreTrainedModel, GenerationMixin):
 
         for grad_step in range(num_steps_with_grad):
             xk = x
-            x, block_idx = self.core_block_forward(
+            x, block_idx = self._maybe_checkpoint_core_block(
                 xk, input_embeds, freqs_cis, mask, past_key_values, block_idx, num_steps_no_grad + grad_step
             )
         return self.transformer.ln_f(x), num_steps_no_grad, num_steps_with_grad, xk.detach(), block_idx  # type: ignore # types broken in 2.6+
@@ -763,13 +768,73 @@ class RavenForCausalLM(RavenPreTrainedModel, GenerationMixin):
         block_idx: torch.Tensor,
         current_step: int | Tensor,
     ):
+        block_idx = block_idx.detach().clone()  # line only included to convince torch.checkpointing
         x = self._maybe_inject_noise(x, current_step)
         x = self.transformer.adapter(torch.cat([x, input_embeds.to(x.device)], dim=-1))  # type: ignore # types broken in 2.6+
         for block in self.transformer.core_block:  # type: ignore # types broken in 2.6+
             block_idx += 1
             x = block(x, freqs_cis, block_idx, mask, past_key_values)
+
         return x, block_idx
 
+    @torch._dynamo.disable(recursive=False)  # type: ignore
+    def randomized_iteration_sampler(self) -> tuple[torch.Tensor, torch.Tensor]:
+        """Outputs are long tensors so that they can be passed through compiled functions"""
+        t = max(self.config.mean_recurrence - self.config.mean_backprop_depth, 0)
+        s = self.config.mean_backprop_depth
+        if torch.rand((1,)).is_meta:  # annoying clause to make meta-tensor-based flop counting work
+            # these values are only the mean TFLOPs of the randomized sampler
+            # Note that this clause also breaks the contract, and returns ints in meta tensor mode
+            return t, s  # type: ignore
+        if self.training:
+            sigma = 0.5
+            mu = math.log(t + s) - (sigma**2 / 2)
+            rate = torch.zeros((1,)).log_normal_(mean=mu, std=sigma)
+            p = torch.poisson(torch.tensor([rate], dtype=torch.float)) + 1
+            n = torch.clamp(p - s, min=0)
+            k = torch.as_tensor(torch.minimum(torch.as_tensor(s), p))
+        else:
+            n, k = torch.as_tensor(self.config.mean_recurrence), torch.as_tensor(0)
+
+        return n.to(dtype=torch.long), k.to(dtype=torch.long)
+
+    def initialize_state(self, input_embeds, scale: float = 1.0):
+        x = torch.randn_like(input_embeds)
+        std = self.config.init_values["std"] * scale
+        if std > 0:
+            torch.nn.init.trunc_normal_(x, mean=0.0, std=std, a=-3 * std, b=3 * std)
+            if self.emb_scale != 1:
+                x = x * self.emb_scale
+        else:
+            x.zero_()
+        return x
+
+    def _maybe_inject_noise(self, x, current_step, renorm=True):
+        if self.config.test_time_noise > 0:
+            n = self.config.test_time_noise * self.config.init_values["std"] * self.emb_scale
+            if self.config.test_time_noise_type == "geom":
+                step1 = torch.as_tensor(current_step + 1, device=x.device)  # need to cast for compile
+                x = x * (1 - n / step1) + torch.randn_like(x) * n / step1
+            elif self.config.test_time_noise_type == "sqrt":
+                step1sqrt = torch.as_tensor(current_step + 1, device=x.device).sqrt()  # need to cast for compile
+                x = x * (1 - n / step1sqrt) + torch.randn_like(x) * n / step1sqrt
+            elif self.config.test_time_noise_type == "line":
+                noise = max(n, (self.config.mean_recurrence - current_step) / self.config.mean_recurrence)  # type: ignore
+                x = x * (1 - noise) + torch.randn_like(x) * noise
+            elif self.config.test_time_noise_type == "chi":
+                noise = 2 * torch.rand(1, device=x.device, dtype=x.dtype) * n
+                x = x * (1 - noise) + torch.randn_like(x) * noise
+            elif self.config.test_time_noise_type == "fixed":
+                x = x * (1 - n) + torch.randn_like(x) * n
+            else:
+                raise ValueError()
+
+            if renorm:
+                x = self.transformer.core_block[-1].norm_4(x)  # type: ignore moduledict types still broken in pytorch
+        return x
+
+    """ ------------------ Alternative interfaces into the model forward ---------------------------------------- """
+
     @torch.no_grad()
     def iterate_one_step(
         self,
@@ -865,61 +930,135 @@ class RavenForCausalLM(RavenPreTrainedModel, GenerationMixin):
             input_embeds = block(input_embeds, freqs_cis, block_idx, prepared_attn_mask, past_key_values)
         return input_embeds, block_idx
 
-    @torch._dynamo.disable(recursive=False)  # type: ignore
-    def randomized_iteration_sampler(self) -> tuple[torch.Tensor, torch.Tensor]:
-        """Outputs are long tensors so that they can be passed through compiled functions"""
-        t = max(self.config.mean_recurrence - self.config.mean_backprop_depth, 0)
-        s = self.config.mean_backprop_depth
-        if torch.rand((1,)).is_meta:  # annoying clause to make meta-tensor-based flop counting work
-            # these values are only the mean TFLOPs of the randomized sampler
-            # Note that this clause also breaks the contract, and returns ints in meta tensor mode
-            return t, s  # type: ignore
-        if self.training:
-            sigma = 0.5
-            mu = math.log(t + s) - (sigma**2 / 2)
-            rate = torch.zeros((1,)).log_normal_(mean=mu, std=sigma)
-            p = torch.poisson(torch.tensor([rate], dtype=torch.float)) + 1
-            n = torch.clamp(p - s, min=0)
-            k = torch.as_tensor(torch.minimum(torch.as_tensor(s), p))
-        else:
-            n, k = torch.as_tensor(self.config.mean_recurrence), torch.as_tensor(0)
+    @torch.no_grad()
+    def _prefill_with_varied_exit_steps(
+        self,
+        input_ids: torch.Tensor,
+        exit_evaluator: "PerIterationExitEvaluator",
+        past_key_values: Optional[ValidCache] = None,
+        init_scale: float = 1.0,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, ValidCache, List[int]]:
+        """ "
+        Note that this the opposite of a real prefill, it goes token-by token and can adaptively exit on each.
+        Use for scientific experiments.
+        """
+        # currently the cache doesn't support batching with adaptive compute
+        assert input_ids.shape[0] == 1
 
-        return n.to(dtype=torch.long), k.to(dtype=torch.long)
+        if past_key_values is None:
+            past_key_values = HuginnDynamicCache()
+        attention_mask = None
+        output = torch.empty(
+            (input_ids.shape[0], 0, self.config.vocab_size), device=input_ids.device, dtype=torch.float
+        )
+        compute_steps = []
+        for pos in range(input_ids.shape[1]):
+            aux_inputs = {
+                "cache_position": pos,
+                "past_key_values": past_key_values,
+                "attention_mask": attention_mask,
+            }
+            freqs_cis = self.freqs_cis[:, pos]
+            embedded_inputs, block_idx = self.embed_inputs(input_ids[:, pos].unsqueeze(1), **aux_inputs)
 
-    def initialize_state(self, input_embeds, scale: float = 1.0):
-        x = torch.randn_like(input_embeds)
-        std = self.config.init_values["std"] * scale
-        if std > 0:
-            torch.nn.init.trunc_normal_(x, mean=0.0, std=std, a=-3 * std, b=3 * std)
-            if self.emb_scale != 1:
-                x = x * self.emb_scale
+            current_latents = self.initialize_state(embedded_inputs, scale=init_scale)
+            exit_evaluator.init(current_latents)
+
+            # Main recurrence
+            for compute_step in range(self.config.mean_recurrence):
+                current_latents, block_idx, _ = self.iterate_one_step(
+                    embedded_inputs,
+                    current_latents,
+                    block_idx=block_idx,
+                    **aux_inputs,
+                    current_step=compute_step,
+                )
+                new_exits, _, _ = exit_evaluator.check(self, current_latents, aux_inputs)
+                if new_exits.any():
+                    break
+            compute_steps.append(compute_step + 1)
+
+            x = self.transformer.ln_f(current_latents)  # type: ignore
+
+            # Coda layers
+            block_idx = torch.tensor(0, device=torch.device("cpu"), dtype=torch.long)  # use negative indices for head
+            for block in self.transformer.coda:  # type: ignore # types broken in 2.6+
+                block_idx -= 1
+                x = block(x, freqs_cis, block_idx, attention_mask, past_key_values)
+
+            x = self.transformer.ln_f(x)  # type: ignore
+            logits = self.lm_head(x).float()
+            output = torch.cat([output, logits], dim=1)
+        return output, past_key_values, compute_steps  # type: ignore
+
+    @torch.no_grad()
+    def forward_with_adaptive_compute(
+        self,
+        input_ids: torch.Tensor,
+        exit_evaluator: "PerIterationExitEvaluator",
+        labels: Optional[torch.Tensor] = None,
+        past_key_values: Optional[ValidCache] = None,
+        output_details: dict = {
+            "return_logits": True,
+            "return_latents": True,
+            "return_head": False,
+            "return_stats": False,
+        },
+        init_scale: float = 1.0,
+        **kwargs,
+    ) -> CausalLMOutputRecurrentLatents:
+        """This forward call does not make use of the causal nature of transformers, it runs token-by token!
+        Do not use this function for anything other than scientific experiments with adaptive compute!
+        """
+        logits, past_key_values, compute_steps = self._prefill_with_varied_exit_steps(
+            input_ids, exit_evaluator, past_key_values, init_scale
+        )
+        if labels is not None:
+            loss = torch.nn.functional.cross_entropy(logits.view(-1, logits.shape[-1]), labels.view(-1))
+            log_ppl = loss.clone().detach()
         else:
-            x.zero_()
-        return x
+            loss, log_ppl = torch.as_tensor(0.0), torch.as_tensor(0.0)
 
-    def _maybe_inject_noise(self, x, current_step, renorm=False):
-        if self.config.test_time_noise > 0:
-            n = self.config.test_time_noise * self.config.init_values["std"] * self.emb_scale
-            if self.config.test_time_noise_type == "geom":
-                step1 = torch.as_tensor(current_step + 1, device=x.device)  # need to cast for compile
-                x = x * (1 - n / step1) + torch.randn_like(x) * n / step1
-            elif self.config.test_time_noise_type == "sqrt":
-                step1sqrt = torch.as_tensor(current_step + 1, device=x.device).sqrt()  # need to cast for compile
-                x = x * (1 - n / step1sqrt) + torch.randn_like(x) * n / step1sqrt
-            elif self.config.test_time_noise_type == "line":
-                noise = max(n, (self.config.mean_recurrence - current_step) / self.config.mean_recurrence)  # type: ignore
-                x = x * (1 - noise) + torch.randn_like(x) * noise
-            elif self.config.test_time_noise_type == "chi":
-                noise = 2 * torch.rand(1, device=x.device, dtype=x.dtype) * n
-                x = x * (1 - noise) + torch.randn_like(x) * noise
-            elif self.config.test_time_noise_type == "fixed":
-                x = x * (1 - n) + torch.randn_like(x) * n
-            else:
-                raise ValueError()
+        return CausalLMOutputRecurrentLatents(
+            loss=loss,
+            log_ppl=log_ppl,
+            logits=logits if output_details["return_logits"] else None,
+            past_key_values=None,
+            hidden_states=None,
+            latent_states=None,
+            attention_maps=None,
+            stats={"compute_steps": compute_steps},
+        )
 
-        if renorm:
-            x = self.transformer.core_block[-1].norm_4(x)  # type: ignore moduledict types still broken in pytorch
-        return x
+    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)
+        residual_diff = (x - latent_states).norm(dim=-1)
+        rel_residual = residual_diff / latent_states.norm(dim=-1)
+        stats = {
+            "entropy": prob_entropy,
+            "residual_diff": residual_diff,
+            "rel_residual": rel_residual,
+            "num_steps_no_grad": num_steps_no_grad,
+            "num_steps_with_grad": num_steps_with_grad,
+        }
+        return stats
+
+    def _maybe_checkpoint_core_block(self, *args, **kwargs) -> tuple[Tensor, Tensor]:
+        if self.gradient_checkpointing:
+            return checkpoint(
+                self.core_block_forward,
+                *args,
+                use_reentrant=False,
+                preserve_rng_state=False,
+                determinism_check="none",
+                **kwargs,
+            )  # type: ignore
+        else:
+            return self.core_block_forward(*args)
+
+    """"------------------------------------------Generation Utilities from here----------------------------------"""
 
     def prepare_inputs_for_generation(
         self,
@@ -971,11 +1110,11 @@ class RavenForCausalLM(RavenPreTrainedModel, GenerationMixin):
     def generate(self, *args, **kwargs):
         """Dispatcher - use HF generate in all normal cases."""
         self.generation_config = args[1] if len(args) > 1 else self.generation_config
-        if any(k in kwargs for k in ("criterion", "exit_threshold")):
-            # print("Dispatching to custom generate_adaptive function call")
+        if any(k in kwargs for k in ("criterion", "exit_threshold", "exit_evaluator")):
             return self.generate_with_adaptive_compute(*args, **kwargs)
+        elif any(k in kwargs for k in ("draft_steps", "lookahead_for_draft", "verification_threshold")):
+            return self.generate_speculative(*args, **kwargs)
         elif "continuous_compute" in kwargs:
-            # print("Dispatching to custom generate_minimal function call")
             return self.generate_minimal(*args, **kwargs)
         else:
             return super().generate(*args, **kwargs)
@@ -1013,7 +1152,7 @@ class RavenForCausalLM(RavenPreTrainedModel, GenerationMixin):
                 lookup_strategy=cache_lookup_strategy,
             )
         model_kwargs["use_cache"] = True
-        model_kwargs = self._get_initial_cache_position(input_ids, model_kwargs)
+        model_kwargs = self._get_initial_cache_position(input_ids.shape[1], input_ids.device, model_kwargs)
         return model_kwargs, generation_config, max_new_tokens
 
     @torch.no_grad()
@@ -1030,7 +1169,7 @@ class RavenForCausalLM(RavenPreTrainedModel, GenerationMixin):
     ) -> Union[torch.Tensor, dict[str, Any]]:
         """Minimal single-sequence generation. Template for more complicated generate tasks"""
         model_kwargs, generation_config, max_new_tokens = self._prep_generate_args(
-            input_ids, generation_config, cache_lookup_strategy
+            input_ids, generation_config, cache_lookup_strategy, model_kwargs
         )
         stop_tokens = self._get_stops(generation_config, tokenizer, model_kwargs).to(input_ids.device)
         unfinished_sequences = torch.ones(input_ids.shape[0], dtype=torch.long, device=input_ids.device)
@@ -1093,15 +1232,25 @@ class RavenForCausalLM(RavenPreTrainedModel, GenerationMixin):
         tokenizer=None,
         streamer=None,
         continuous_compute=False,  # warm-start state / continuous CoT
-        criterion="none",  # off by default, turn on by choosing an exit criterion
+        criterion="none",  # adaptive compute is off by default, turn on by choosing an exit criterion
         exit_threshold: Union[str, float, int] = "auto",
         init_scale: float = 1.0,
         cache_lookup_strategy: str = "full",
+        do_not_exit_in_prefill: bool = False,
+        min_steps: int = 0,
+        check_criterion_every_n_steps=1,
+        exit_evaluator: "Optional[PerIterationExitEvaluator]" = None,  # optional plugin of a new exit eval object
         **model_kwargs,
     ) -> Union[torch.Tensor, GenerateDecoderOnlyOutput]:
         """
         Generate tokens with adaptive compute. This is NOT the most efficient implementation.
         For batches, on each token, we iterate until the entire batch finishes.
+        Note: While the method can be used batched, and will produce sensible results, this cannot be used to evaluate
+              the success of adaptive compute methods, which should only ever be benchmarked with batch_size=1.
+              This is because the KV-cache entries are necessarily batched and so contain entries equal to the sequence
+              with the largest number of steps in the whole batch, and these KV states, which would not have been computed
+              if there was only one (short compute) sequence in the batch, will be picked up by later compute steps,
+              making early exits look better than they are.
         """
         model_kwargs, generation_config, max_new_tokens = self._prep_generate_args(
             input_ids, generation_config, cache_lookup_strategy, model_kwargs
@@ -1120,8 +1269,11 @@ class RavenForCausalLM(RavenPreTrainedModel, GenerationMixin):
         # Track which sequences have finished (using unfinished_sequences to match generate_minimal)
         unfinished_sequences = torch.ones(batch_size, dtype=torch.long, device=input_ids.device)
 
+        if exit_evaluator is None:
+            exit_evaluator = get_adaptive_exit_evaluator(self, criterion, exit_threshold)
+
         # Generate tokens
-        for _ in range(max_new_tokens):
+        for token_step_in_sequence in range(max_new_tokens):
             # Adaptive compute forward
             model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
             aux_inputs = {
@@ -1134,38 +1286,20 @@ class RavenForCausalLM(RavenPreTrainedModel, GenerationMixin):
                 else model_kwargs["input_states"]
             )
 
+            # Initialize next_states for continuous compute
+            if continuous_compute:
+                next_states = current_latents[:, -1:, :].clone()
+
             # Initialize criterion tracking for each sequence in batch
             exit_values_per_seq = [[] for _ in range(batch_size)]
             compute_steps_per_seq = [0] * batch_size
             exit_reached = torch.zeros(batch_size, dtype=torch.bool, device=input_ids.device)
 
-            # Set up criterions based on selected strategy
-            if criterion == "entropy-diff":
-                entropy = torch.ones(batch_size, device=input_ids.device) * 100.0
-                exit_threshold = 1e-3 if exit_threshold == "auto" else float(exit_threshold)
-            elif criterion == "latent-diff":
-                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(batch_size, V, dtype=torch.float, 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 = torch.zeros(batch_size, dtype=torch.long, device=input_ids.device)
-                current_argmax = torch.ones(batch_size, dtype=torch.long, device=input_ids.device) * -1
-                exit_threshold = 5 if exit_threshold == "auto" else int(exit_threshold)
-            elif criterion == "none":
-                exit_threshold = 1.0 if exit_threshold == "auto" else float(exit_threshold)
-            else:
-                raise ValueError("Invalid adaptive compute strategy.")
-
-            next_token_logits = None
+            outputs, next_token_logits = None, None
+            exit_evaluator.init(current_latents)
 
             # Iterate through compute steps
             for compute_step in range(max_steps):
-                prev_latents = current_latents.clone()
                 current_latents, block_idx, _ = self.iterate_one_step(
                     embedded_inputs,
                     current_latents,
@@ -1174,94 +1308,70 @@ class RavenForCausalLM(RavenPreTrainedModel, GenerationMixin):
                     current_step=compute_step,
                 )
 
-                if _ > 0:  # do not exit in prefill
-                    # Check exit condition for each sequence in batch
-                    if criterion == "entropy-diff":
-                        prev_entropy = entropy
-                        outputs = self.predict_from_latents(current_latents, **aux_inputs)
-                        logits: torch.Tensor = outputs.logits  # type: ignore
-                        probs = F.softmax(logits[:, -1, :], dim=-1)
-                        entropy = -torch.sum(probs * torch.log(probs + 1e-10), dim=-1)
-                        exit_values = (entropy - prev_entropy).abs()
-                    elif criterion == "latent-diff":
-                        norm_diff = (prev_latents - current_latents).norm(dim=-1) / current_latents.norm(dim=-1)
-                        exit_values = norm_diff.mean(dim=-1)
-                    elif "kl" in criterion:
-                        outputs = self.predict_from_latents(current_latents, **aux_inputs)
-                        logits: torch.Tensor = outputs.logits  # type: ignore
-                        prev_log_probs = log_probs
-                        if criterion == "minp-kl":
-                            probs = F.softmax(logits[:, -1, :].float(), dim=-1)
-                            max_probs = probs.max(dim=-1, keepdim=True)[0]
-                            probs_mask = probs < (0.1 * max_probs)
-                            masked_probs = probs.clone()
-                            masked_probs[probs_mask] = 1 / V
-                            probs = masked_probs / masked_probs.sum(dim=-1, keepdim=True)
-                            log_probs = probs.log()
-                        else:
-                            log_probs = F.log_softmax(logits[:, -1, :].float(), dim=-1)
-                        exit_values = F.kl_div(log_probs, prev_log_probs, reduction="none", log_target=True).sum(dim=-1)
-                    elif criterion == "argmax-stability":
-                        prev_argmax = current_argmax
+                # Skip checking exit conditions if min_steps not met, or not checking this step, or in prefill
+                if (
+                    compute_step < min_steps
+                    or (compute_step - min_steps) % check_criterion_every_n_steps != 0
+                    or (do_not_exit_in_prefill and token_step_in_sequence == 0)
+                ):
+                    continue
+
+                # Otherwise check for new exits, potentially by evaluating the coda:
+                new_exits, outputs, exit_values = exit_evaluator.check(self, current_latents, aux_inputs)
+
+                # Record values and check exits for each sequence
+                for i in range(batch_size):
+                    if not exit_reached[i] and unfinished_sequences[i].bool():
+                        exit_values_per_seq[i].append(exit_values[i].item())
+
+                new_exits = new_exits & ~exit_reached & unfinished_sequences.bool()
+
+                if new_exits.any():
+                    exit_reached = exit_reached | new_exits
+                    if outputs is not None:
+                        logits = outputs.logits
+                    else:
+                        # For latent-based criteria, compute outputs when we need them
                         outputs = self.predict_from_latents(current_latents, **aux_inputs)
-                        logits: torch.Tensor = outputs.logits  # type: ignore
-                        current_argmax = logits[:, -1, :].argmax(dim=-1)
-                        stable_for_n_steps = torch.where(
-                            current_argmax == prev_argmax, stable_for_n_steps + 1, torch.zeros_like(stable_for_n_steps)
-                        )
-                        exit_values = stable_for_n_steps
-                    elif criterion == "none":
-                        exit_values = torch.ones(batch_size, device=input_ids.device) * 2.0 * exit_threshold
-
-                    # Record values and check exits for each sequence
+                        logits = outputs.logits
+
+                    if next_token_logits is None:
+                        next_token_logits = logits[:, -1, :].to(**logit_type)  # type: ignore
+                    else:
+                        next_token_logits[new_exits] = logits[new_exits, -1, :].to(**logit_type)  # type: ignore
+
                     for i in range(batch_size):
-                        if not exit_reached[i] and unfinished_sequences[i].bool():
-                            exit_values_per_seq[i].append(exit_values[i].item())
+                        if new_exits[i]:
+                            compute_steps_per_seq[i] = compute_step + 1
 
-                    # Check for new exits, respecting unfinished_sequences
-                    new_exits = (
-                        exit_values < exit_threshold
-                        if criterion != "argmax-stability"
-                        else exit_values >= exit_threshold
-                    )
-                    new_exits = new_exits & ~exit_reached & unfinished_sequences.bool()
-
-                    if new_exits.any():
-                        exit_reached = exit_reached | new_exits
-                        if criterion == "latent-diff":
-                            # Normally we don't compute the output for latent-diff, but when there is an exit,
-                            # we need to compute and save the output
-                            outputs = self.predict_from_latents(current_latents, **aux_inputs)
-                            logits: torch.Tensor = outputs.logits  # type: ignore
-                        if next_token_logits is None:
-                            next_token_logits = logits[:, -1, :].to(**logit_type)  # type: ignore
-                        else:
-                            for i in range(batch_size):
-                                if new_exits[i]:
-                                    next_token_logits[i] = logits[i, -1, :].to(**logit_type)  # type: ignore
-                        for i in range(batch_size):
-                            if new_exits[i]:
-                                compute_steps_per_seq[i] = compute_step + 1
-
-                    # If all sequences have exited or finished, break early
-                    if (exit_reached | ~unfinished_sequences.bool()).all():
-                        break
-            # This else is if the for loop finished without breaking
+                    # Update continuous compute states for newly exited sequences
+                    if continuous_compute:
+                        next_states[new_exits] = current_latents[new_exits, -1:, :]
+
+                # If all sequences have exited or finished, break early
+                if (exit_reached | ~unfinished_sequences.bool()).all():
+                    break
+
+            # This else triggers if the for loop finishes without breaking:
             else:
-                outputs = self.predict_from_latents(current_latents, **aux_inputs)
+                if outputs is None:
+                    outputs = self.predict_from_latents(current_latents, **aux_inputs)
 
                 # For sequences that didn't exit early, use the final logits
                 if next_token_logits is None:
                     next_token_logits = outputs.logits[:, -1, :].to(**logit_type)  # type: ignore
+                    for i in range(batch_size):
+                        compute_steps_per_seq[i] = max_steps
                 else:
                     for i in range(batch_size):
                         if not exit_reached[i] and unfinished_sequences[i].bool():
                             next_token_logits[i] = outputs.logits[i, -1, :].to(**logit_type)  # type: ignore
                             compute_steps_per_seq[i] = max_steps
-
             # Save latent states for continuous compute if enabled
             if continuous_compute:
-                model_kwargs["input_states"] = current_latents[:, -1:, :]
+                still_running = ~exit_reached & unfinished_sequences.bool()
+                next_states[still_running] = current_latents[still_running, -1:, :]
+                model_kwargs["input_states"] = next_states
 
             # Record compute steps for this token generation
             compute_steps.append([compute_steps_per_seq, exit_values_per_seq])
@@ -1276,7 +1386,7 @@ class RavenForCausalLM(RavenPreTrainedModel, GenerationMixin):
                 streamer.put(next_token.cpu())
 
             # Update model kwargs for next iteration
-            model_kwargs = self._update_model_kwargs_for_generation(outputs, model_kwargs)
+            model_kwargs = self._update_model_kwargs_for_generation(outputs, model_kwargs)  # type: ignore
 
             # Check for stop tokens and update unfinished sequences
             for i in range(batch_size):
@@ -1309,62 +1419,6 @@ class RavenForCausalLM(RavenPreTrainedModel, GenerationMixin):
             )
         return input_ids
 
-    def _get_stops(self, generation_config, tokenizer, model_kwargs):
-        stop_tokens = {65504, 65505, 65508}  # begin_text, end_text, end_turn
-        if generation_config.eos_token_id is not None:
-            stop_tokens.add(generation_config.eos_token_id)
-        if "stopping_criteria" in model_kwargs and tokenizer is None:
-            tokenizer = model_kwargs["stopping_criteria"][0].tokenizer
-        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 _sample_next_token(self, next_token_logits, generation_config):
-        """Helper function to sample the next token."""
-        if generation_config.do_sample:
-            if generation_config.temperature:
-                next_token_logits = next_token_logits.float() / generation_config.temperature
-
-            probs = F.softmax(next_token_logits, dim=-1)
-
-            # Apply top_k
-            if generation_config.top_k:
-                top_k_values, _ = torch.topk(probs, generation_config.top_k, dim=-1)
-                min_values = top_k_values[:, -1].unsqueeze(-1).expand_as(probs)
-                probs = torch.where(probs < min_values, torch.zeros_like(probs), probs)
-
-            # Apply top_p (nucleus sampling)
-            if generation_config.top_p:
-                sorted_probs, sorted_indices = torch.sort(probs, descending=True, dim=-1)
-                cumulative_probs = torch.cumsum(sorted_probs, dim=-1)
-
-                # Create mask for probs to keep
-                remove_indices = cumulative_probs > generation_config.top_p
-                remove_indices[:, 0] = False  # Keep at least the top probability
-
-                # Convert sorted indices mask back to original indices mask
-                mask = torch.zeros_like(probs, dtype=torch.bool)
-                for i in range(probs.shape[0]):
-                    mask[i, sorted_indices[i, remove_indices[i]]] = True
-
-                probs = torch.where(mask, torch.zeros_like(probs), probs)
-
-            # Apply min_p
-            if generation_config.min_p:
-                max_probs = probs.max(dim=-1, keepdim=True)[0]
-                min_p_threshold = generation_config.min_p * max_probs
-                probs = torch.where(probs < min_p_threshold, torch.zeros_like(probs), probs)
-
-            # Renormalize probabilities
-            probs = probs / probs.sum(dim=-1, keepdim=True).clamp(min=1e-10)
-
-            # Sample from the distribution
-            return torch.multinomial(probs, num_samples=1)
-        else:
-            return torch.argmax(next_token_logits, dim=-1, keepdim=True)
-
     @torch.no_grad()
     def generate_speculative(
         self,
@@ -1546,22 +1600,69 @@ class RavenForCausalLM(RavenPreTrainedModel, GenerationMixin):
             )
         return input_ids
 
-    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)
-        residual_diff = (x - latent_states).norm(dim=-1)
-        rel_residual = residual_diff / latent_states.norm(dim=-1)
-        stats = {
-            "entropy": prob_entropy,
-            "residual_diff": residual_diff,
-            "rel_residual": rel_residual,
-            "num_steps_no_grad": num_steps_no_grad,
-            "num_steps_with_grad": num_steps_with_grad,
-        }
-        return stats
+    def _get_stops(self, generation_config, tokenizer, model_kwargs):
+        stop_tokens = {65504, 65505, 65508}  # begin_text, end_text, end_turn
+        if generation_config.eos_token_id is not None:
+            try:
+                stop_tokens.update(generation_config.eos_token_id)
+            except TypeError:
+                stop_tokens.add(generation_config.eos_token_id)
+        if "stopping_criteria" in model_kwargs and tokenizer is None:
+            tokenizer = model_kwargs["stopping_criteria"][0].tokenizer
+        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 _sample_next_token(self, next_token_logits, generation_config):
+        """Helper function to sample the next token."""
+        if generation_config.do_sample:
+            if generation_config.temperature:
+                next_token_logits = next_token_logits.float() / generation_config.temperature
+
+            probs = F.softmax(next_token_logits, dim=-1)
+
+            # Apply top_k
+            if generation_config.top_k:
+                top_k_values, _ = torch.topk(probs, generation_config.top_k, dim=-1)
+                min_values = top_k_values[:, -1].unsqueeze(-1).expand_as(probs)
+                probs = torch.where(probs < min_values, torch.zeros_like(probs), probs)
+
+            # Apply top_p (nucleus sampling)
+            if generation_config.top_p:
+                sorted_probs, sorted_indices = torch.sort(probs, descending=True, dim=-1)
+                cumulative_probs = torch.cumsum(sorted_probs, dim=-1)
+
+                # Create mask for probs to keep
+                remove_indices = cumulative_probs > generation_config.top_p
+                remove_indices[:, 0] = False  # Keep at least the top probability
+
+                # Convert sorted indices mask back to original indices mask
+                mask = torch.zeros_like(probs, dtype=torch.bool)
+                for i in range(probs.shape[0]):
+                    mask[i, sorted_indices[i, remove_indices[i]]] = True
+
+                probs = torch.where(mask, torch.zeros_like(probs), probs)
+
+            # Apply min_p
+            if generation_config.min_p:
+                max_probs = probs.max(dim=-1, keepdim=True)[0]
+                min_p_threshold = generation_config.min_p * max_probs
+                probs = torch.where(probs < min_p_threshold, torch.zeros_like(probs), probs)
+
+            # Renormalize probabilities
+            probs = probs / probs.sum(dim=-1, keepdim=True).clamp(min=1e-10)
+
+            # Sample from the distribution
+            return torch.multinomial(probs, num_samples=1)
+        else:
+            return torch.argmax(next_token_logits, dim=-1, keepdim=True)
+
+
+################################ Model Utils #######################################################################
 
 
-#################################### Utils #######################################################################
 def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0, condense_ratio: int = 1):
     with torch.autocast("cuda", enabled=False):
         inv_freqs = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
@@ -1587,6 +1688,204 @@ def apply_rotary_emb_complex_like(q: Tensor, k: Tensor, freqs_cis: Tensor) -> tu
         return torch.split(rotated_qk.type_as(q), q.shape[2], dim=2)  # type: ignore
 
 
+#################################### Adaptive Compute Exit Evaluators ##########################################
+
+Exit = Tuple[torch.Tensor, Optional[CausalLMOutputRecurrentLatents], torch.Tensor]
+
+
+class PerIterationExitEvaluator:
+    """Base class for exit evaluators that check after each recurrent step."""
+
+    def init(self, initial_latents: torch.Tensor):
+        """Initialize evaluator state."""
+
+    def check(self, model: "RavenForCausalLM", latents: torch.Tensor, aux_inputs: dict) -> Exit:
+        """Returns (should_exit, outputs (or None), exit_values)"""
+        raise NotImplementedError()
+
+
+class NoOpExitEvaluator(PerIterationExitEvaluator):
+    """Exit evaluator that never exits early."""
+
+    def check(self, model: "RavenForCausalLM", latents: torch.Tensor, aux_inputs: dict) -> Exit:
+        return (
+            torch.zeros(latents.shape[0], device=latents.device, dtype=torch.bool),
+            None,
+            torch.zeros(latents.shape[0], device=latents.device),
+        )
+
+
+class EntropyDiffExitEvaluator(PerIterationExitEvaluator):
+    """Exit based on change in output entropy."""
+
+    def __init__(self, exit_threshold: Union[str, float] = "auto"):
+        self.exit_threshold = 1e-3 if exit_threshold == "auto" else float(exit_threshold)
+
+    def init(self, initial_latents: torch.Tensor):
+        batch_size = initial_latents.shape[0]
+        self.prev_entropy = torch.ones(batch_size, device=initial_latents.device) * 100.0
+
+    def check(self, model: "RavenForCausalLM", latents: torch.Tensor, aux_inputs: dict) -> Exit:
+        outputs = model.predict_from_latents(latents, **aux_inputs)
+        logits: torch.Tensor = outputs.logits  # type: ignore
+        probs = F.softmax(logits[:, -1, :], dim=-1)
+        entropy = -torch.sum(probs * torch.log(probs + 1e-10), dim=-1)
+        exit_values = (entropy - self.prev_entropy).abs()
+        self.prev_entropy = entropy
+        return exit_values < self.exit_threshold, outputs, exit_values
+
+
+class LatentDiffExitEvaluator(PerIterationExitEvaluator):
+    """Exit based on change in latent states."""
+
+    def __init__(self, exit_threshold: Union[str, float] = "auto"):
+        self.exit_threshold = 0.03 if exit_threshold == "auto" else float(exit_threshold)
+
+    def init(self, initial_latents: torch.Tensor):
+        self.prev_latents = initial_latents.clone().detach()
+
+    def check(self, model: "RavenForCausalLM", latents: torch.Tensor, aux_inputs: dict) -> Exit:
+        exit_values = ((latents - self.prev_latents).norm(dim=-1) / latents.norm(dim=-1)).mean(dim=-1)
+        self.prev_latents = latents.clone().detach()
+        return exit_values < self.exit_threshold, None, exit_values
+
+
+class KLExitEvaluator(PerIterationExitEvaluator):
+    """Exit based on KL divergence between successive outputs."""
+
+    def __init__(self, model: "RavenForCausalLM", exit_threshold: Union[str, float] = "auto"):
+        self.exit_threshold = 0.001 if exit_threshold == "auto" else float(exit_threshold)
+        self.V = model.config.padded_vocab_size
+
+    def init(self, initial_latents: torch.Tensor):
+        batch_size = initial_latents.shape[0]
+        self.prev_log_probs = ((1 / self.V) * torch.ones(batch_size, self.V, device=initial_latents.device)).log()
+
+    def check(self, model: "RavenForCausalLM", latents: torch.Tensor, aux_inputs: dict) -> Exit:
+        outputs = model.predict_from_latents(latents, **aux_inputs)
+        logits: torch.Tensor = outputs.logits  # type: ignore
+        log_probs = F.log_softmax(logits[:, -1, :].float(), dim=-1)
+        exit_values = F.kl_div(log_probs, self.prev_log_probs, reduction="none", log_target=True).sum(dim=-1)
+        self.prev_log_probs = log_probs
+        return exit_values < self.exit_threshold, outputs, exit_values
+
+
+class MinKLExitEvaluator(PerIterationExitEvaluator):
+    """Exit based on min-p filtered KL divergence."""
+
+    def __init__(self, model: "RavenForCausalLM", exit_threshold: Union[str, float] = "auto"):
+        self.exit_threshold = 1e-5 if exit_threshold == "auto" else float(exit_threshold)
+        self.V = model.config.padded_vocab_size
+
+    def init(self, initial_latents: torch.Tensor):
+        batch_size = initial_latents.shape[0]
+        self.prev_log_probs = ((1 / self.V) * torch.ones(batch_size, self.V, device=initial_latents.device)).log()
+
+    def _calc_minp_log_probs(self, logits: torch.Tensor) -> torch.Tensor:
+        probs = F.softmax(logits[:, -1, :], dim=-1)
+        max_probs = probs.max(dim=-1, keepdim=True)[0]
+        probs_mask = probs < (0.1 * max_probs)
+        masked_probs = probs
+        masked_probs[probs_mask] = 1 / self.V
+        probs = masked_probs / masked_probs.sum(dim=-1, keepdim=True)
+        return probs.log()
+
+    def check(self, model: "RavenForCausalLM", latents: torch.Tensor, aux_inputs: dict) -> Exit:
+        outputs = model.predict_from_latents(latents, **aux_inputs)
+        logits: torch.Tensor = outputs.logits  # type: ignore
+        log_probs = self._calc_minp_log_probs(logits)
+        exit_values = F.kl_div(log_probs, self.prev_log_probs, reduction="none", log_target=True).sum(dim=-1)
+        self.prev_log_probs = log_probs
+        return exit_values < self.exit_threshold, outputs, exit_values
+
+
+class ArgmaxStabilityExitEvaluator(PerIterationExitEvaluator):
+    """Exit based on argmax stability over consecutive steps."""
+
+    def __init__(self, exit_threshold: Union[str, int] = "auto"):
+        self.exit_threshold = 5 if exit_threshold == "auto" else int(exit_threshold)
+
+    def init(self, initial_latents: torch.Tensor):
+        batch_size = initial_latents.shape[0]
+        self.prev_argmax = torch.ones(batch_size, dtype=torch.long, device=initial_latents.device) * -1
+        self.stable_for_n_steps = torch.zeros(batch_size, dtype=torch.long, device=initial_latents.device)
+
+    def check(self, model: "RavenForCausalLM", latents: torch.Tensor, aux_inputs: dict) -> Exit:
+        outputs = model.predict_from_latents(latents, **aux_inputs)
+        logits: torch.Tensor = outputs.logits  # type: ignore
+        current_argmax = logits[:, -1, :].argmax(dim=-1)
+        stable_for_n_steps = torch.where(
+            current_argmax == self.prev_argmax, self.stable_for_n_steps + 1, torch.zeros_like(self.stable_for_n_steps)
+        )
+        exit_values = stable_for_n_steps
+        self.prev_argmax = current_argmax
+        self.stable_for_n_steps = stable_for_n_steps
+        return exit_values >= self.exit_threshold, outputs, exit_values
+
+
+class CosineExitEvaluator(PerIterationExitEvaluator):
+    """Exit based on cosine similarity between successive latent states."""
+
+    def __init__(self, exit_threshold: Union[str, float] = "auto"):
+        self.exit_threshold = 1e-3 if exit_threshold == "auto" else float(exit_threshold)
+
+    def init(self, initial_latents: torch.Tensor):
+        self.prev_latents = initial_latents.clone().detach()
+
+    def check(self, model: "RavenForCausalLM", latents: torch.Tensor, aux_inputs: dict) -> Exit:
+        cosine_sim = (
+            (latents * self.prev_latents).sum(dim=-1) / latents.norm(dim=-1) / self.prev_latents.norm(dim=-1)
+        ).mean(dim=1)
+        exit_values = 1 - cosine_sim
+        self.prev_latents = latents.clone().detach()
+        return exit_values < self.exit_threshold, None, exit_values
+
+
+class NumStepsGenerator(PerIterationExitEvaluator):
+    def __init__(self, steps_fn: Callable):
+        self.steps_fn = steps_fn
+        self.counter = 0
+        self.target_steps = 0
+        self.current_step = 0
+
+    def init(self, initial_latents):
+        self.target_steps = self.steps_fn(self.counter)
+        self.counter += 1
+        self.current_step = 0
+
+    def check(self, model: "RavenForCausalLM", latents: torch.Tensor, aux_inputs: dict) -> Exit:
+        self.current_step += 1
+        should_exit = self.current_step >= self.target_steps
+        return (
+            torch.full((latents.shape[0],), should_exit, dtype=torch.bool, device=latents.device),
+            None,
+            torch.zeros(latents.shape[0], device=latents.device),
+        )
+
+
+def get_adaptive_exit_evaluator(
+    model: "RavenForCausalLM", criterion: str, exit_threshold: Union[str, float, int]
+) -> PerIterationExitEvaluator:
+    """Factory function to create appropriate exit evaluator."""
+    if criterion == "entropy-diff":
+        return EntropyDiffExitEvaluator(exit_threshold)
+    elif criterion == "latent-diff":
+        return LatentDiffExitEvaluator(exit_threshold)
+    elif criterion == "cosine":
+        return CosineExitEvaluator(exit_threshold)
+    elif "kl" in criterion:
+        if criterion == "minp-kl":
+            return MinKLExitEvaluator(model, exit_threshold)
+        else:
+            return KLExitEvaluator(model, exit_threshold)
+    elif criterion == "argmax-stability":
+        return ArgmaxStabilityExitEvaluator(exit_threshold)  # type: ignore
+    elif criterion == "none":
+        return NoOpExitEvaluator()
+    else:
+        raise ValueError(f"Invalid adaptive compute strategy: {criterion}")
+
+
 #################################### HF registration ############################################################
 
 from transformers import AutoConfig, AutoModel, AutoModelForCausalLM