[generate] stop synchronizing the accelerator on every decode step

@@ -28,6 +28,7 @@ class CacheLayerMixin(ABC):     """Base, abstract class for a single layer's cache."""      is_compileable = False+    is_croppable = False     supports_early_init = True     # Subclasses can set `_layer_type` to auto-register themselves in the mappings, if the class definition lives in a modeling     # file instead of this file. This allows to update the mapping only when the modeling file is imported, which simplifies imports@@ -117,6 +118,7 @@ class DynamicLayer(CacheLayerMixin):     """      is_sliding = False+    is_croppable = True      def lazy_initialization(self, key_states: torch.Tensor, value_states: torch.Tensor) -> None:         self.dtype, self.device = key_states.dtype, key_states.device@@ -960,6 +962,18 @@ def reorder_cache(self, beam_idx: torch.LongTensor):             if self.is_recurrent_states_initialized[i]:                 self.recurrent_states[i] = self.recurrent_states[i].index_select(0, beam_idx.to(self.device)) +    @property+    def is_croppable(self) -> bool:+        """+        Whether `crop` can put this layer back as it was. This is only supported when there are no recurrent states.+        """+        if any(self.is_recurrent_states_initialized.values()):+            return False+        # If nothing is initialized, return False as we don't yet know whether we will have any recurrent states or no, so let's be+        # extra careful. If a conv states is initialized but no recurrent states are, then we return True as we know that we will never+        # have any recurrent state (they are updated in the same forward)+        return any(self.is_conv_states_initialized.values())+     def activate_past_recording(self):         """         Calling this function will activate past state recording, meaning that a call to `update_conv_states` will@@ -1662,6 +1676,11 @@ def is_initialized(self) -> bool:         layers = [layer for layer in self.layers if layer.supports_early_init]         return len(layers) > 0 and all(layer.is_initialized for layer in layers) +    @property+    def is_croppable(self) -> bool:+        """Whether `crop` can put the whole cache back as it was, so a rollback leaves no trace."""+        return all(layer.is_croppable for layer in self.layers)+     @property     def is_sliding(self) -> list[bool]:         """Return whether the layers of the cache are sliding window"""@@ -2085,6 +2104,10 @@ def is_sliding(self):     def is_compileable(self) -> bool:         return self.self_attention_cache.is_compileable +    @property+    def is_croppable(self) -> bool:+        return self.self_attention_cache.is_croppable+     def activate_past_recording(self):         self.self_attention_cache.activate_past_recording() 

@@ -17,6 +17,7 @@ import inspect import os import warnings+from collections import deque from collections.abc import Callable from contextlib import contextmanager from dataclasses import dataclass@@ -356,6 +357,127 @@ class GenerateBeamEncoderDecoderOutput(ModelOutput): GenerateOutput = GenerateNonBeamOutput | GenerateBeamOutput  +def _undo_generation_steps(num_steps: int, input_ids: torch.LongTensor, *recorded: "tuple | None") -> tuple:+    """+    Undo the last `num_steps` decoding steps, so that they leave no trace in what `generate` returns.+    Note that the cache entries those steps wrote are dropped by `DeferredStopCheck.finish` instead.+    """+    # `[:-0]` is `[:0]`, which would empty everything rather than leave it alone+    if num_steps == 0:+        return (input_ids, *recorded)+    return (input_ids[..., :-num_steps], *(record[:-num_steps] if record else record for record in recorded))+++class StopCheck:+    """+    Decides when the decoding loop should stop, and hands each new token to a streamer.+    """++    def __init__(self, streamer: "BaseStreamer | None" = None):+        self.streamer = streamer++    def __call__(self, unfinished_sequences: torch.Tensor, tokens: torch.Tensor, length: int) -> bool:+        if self.streamer is not None:+            self.streamer.put(tokens.cpu())+        return bool(unfinished_sequences.max() == 0)++    def finish(self) -> int:+        """Flush whatever is still held back, and report how many decoding steps have to be undone."""+        return 0+++class DeferredStopCheck(StopCheck):+    """+    A deferred `StopCheck` that reports whether generation should stop, one step late, so the host never waits on the device.++    Reading `unfinished_sequences.max() == 0` blocks the host until the device has caught up, every step, leaving it unable to queue+    the next step meanwhile. Copying that flag asynchronously and reading it on the *following* step removes the stall. It costs one+    extra `forward` pass, whose results the caller undoes using the count returned by `finish`.+    Streaming needs the tokens themselves on the host, which is the same synchronization, so tokens are streamed one step behind.+    The extra token is never streamed: it is still in flight when the loop breaks, and `finish` drops it.+    """++    def __init__(+        self,+        input_ids: torch.LongTensor,+        max_length: int,+        cache: "Cache | None",+        cache_is_returned: bool,+        streamer: "BaseStreamer | None" = None,+    ):+        super().__init__(streamer)+        self.max_length = max_length+        # We only need to care about rollbacking the cache if we are going to return the Cache, i.e. if the user requested additional+        # outputs or if the user passed an explicit Cache object+        self.cache = cache if cache_is_returned else None+        if self.cache is not None:+            self.cache.activate_past_recording()+        pinned = input_ids.device.type == "cuda"+        self.slots = deque(+            (+                torch.zeros((), dtype=torch.bool, pin_memory=pinned),+                torch.zeros(input_ids.shape[0], dtype=torch.long, pin_memory=pinned),+                torch.Event(device=input_ids.device, blocking=True),+            )+            for _ in range(2)+        )+        self.is_first_step = True+        self.stop_reported = False++    @staticmethod+    def is_supported(device: torch.device, cache: "Cache | None", cache_is_returned: bool, is_assistant: bool) -> bool:+        """+        Whether the stop decision can safely be deferred by a step, in this decoding context.+        In general, we do not defer unless the device is `"mps"`, and if the user requests to return the `Cache`, we need this+        `Cache` to be able to rollback its states correctly, so that the last `forward` can be correctly reverted.+        """+        # Only mps for now, we should enable it for cuda if we observe perf gains - also skip if it's an assistant+        if device.type != "mps" or is_assistant:+            return False+        # Since this is called after prefill, if we still do not have any cache, it means we'll never have one+        if cache is None:+            return True+        # if we don't return the cache, we don't need to bother about activating past recording since it will be dropped anyway+        else:+            return cache.is_croppable or not cache_is_returned++    def __call__(self, unfinished_sequences: torch.Tensor, tokens: torch.Tensor, length: int) -> bool:+        should_stop, tokens_cpu, copy_done = self.slots[0]+        should_stop.copy_(unfinished_sequences.max() == 0, non_blocking=True)+        if self.streamer is not None:+            tokens_cpu.copy_(tokens, non_blocking=True)+        copy_done.record()++        self.slots.rotate()+        should_stop_before, tokens_cpu_before, copy_done_before = self.slots[0]+        copy_done_before.synchronize()+        if not self.is_first_step:+            if self.streamer is not None:+                self.streamer.put(tokens_cpu_before.clone())+            self.stop_reported = bool(should_stop_before)+        self.is_first_step = False+        stopping = self.stop_reported or (self.max_length is not None and length >= self.max_length)+        if not stopping and self.cache is not None:+            self.cache.crop(0)+        return stopping++    def finish(self) -> int:+        for *_, copy_done in self.slots:+            copy_done.synchronize()+        steps_to_undo = 1 if self.stop_reported else 0+        if self.streamer is not None and not steps_to_undo:+            _, tokens_cpu, _ = self.slots[1]+            self.streamer.put(tokens_cpu.clone())+        if self.cache is not None:+            self.cache.crop(-steps_to_undo)+            # We also need to deactivate past_recording, since we are giving the cache back to the user and it may lead to unneeded+            # memory spike on the next prefill+            for layer in self.cache.layers:+                if hasattr(layer, "record_past"):+                    layer.record_past = False+        return steps_to_undo++ class GenerationMixin(ContinuousMixin):     """     A class containing all functions for auto-regressive text generation, to be used as a mixin in model classes.@@ -1955,6 +2077,8 @@ def _prepare_cache_for_generation(                 raise ValueError(                     "Passing a tuple of `past_key_values` is not supported anymore. Please use a `Cache` instance."                 )+            # Marks the cache has user-defined for generate later on+            user_defined_cache._is_user_defined = True             return          # Quick escape route 2: if the user specifies no cache is to be used. (conflicting arguments are handled in@@ -2875,6 +2999,20 @@ def _sample(             is_first_iteration=not generation_config.is_assistant,         ) +        # Decides whether we can defer the stopping criteria to avoid a synchronization point in between every `forward`.+        # Note that it is very important to do this only after the prefill, as we may otherwise call `activate_past_recording` on the+        # Cache, which will cause a huge unneeded memory spike if the prefill is huge and the model would otherwise drop most of the+        # states, such as if it uses sliding window or linear attention+        cache = next((outputs[name] for name in ALL_CACHE_NAMES if name in outputs), None)+        # The cache outlives `generate` if the user asked to return it, or if it is one they passed in.+        cache_is_returned = generation_config.return_dict_in_generate or getattr(cache, "_is_user_defined", False)+        if DeferredStopCheck.is_supported(+            input_ids.device, cache, cache_is_returned, is_assistant=generation_config.is_assistant+        ):+            stop_check = DeferredStopCheck(input_ids, stopping_criteria.max_length, cache, cache_is_returned, streamer)+        else:+            stop_check = StopCheck(streamer)+         with self._optimize_model_for_decode():             while self._has_unfinished_sequences(this_peer_finished, synced_gpus, device=input_ids.device):                 if prefill_consumed:@@ -2933,16 +3071,29 @@ def _sample(                  # update generated ids, model inputs, and length for next step                 input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=-1)-                if streamer is not None:-                    streamer.put(next_tokens.cpu())                  unfinished_sequences = unfinished_sequences & ~stopping_criteria(input_ids, scores)-                this_peer_finished = unfinished_sequences.max() == 0+                this_peer_finished = stop_check(unfinished_sequences, next_tokens, input_ids.shape[1])                  # This is needed to properly delete outputs.logits which may be very large for first iteration                 # Otherwise a reference to outputs is kept which keeps the logits alive in the next iteration                 del outputs +        steps_to_undo = stop_check.finish()+        # We may need to remove the last output if we deferred the stop checks+        if steps_to_undo:+            input_ids, scores, raw_logits, decoder_attentions, cross_attentions, decoder_hidden_states = (+                _undo_generation_steps(+                    steps_to_undo,+                    input_ids,+                    scores,+                    raw_logits,+                    decoder_attentions,+                    cross_attentions,+                    decoder_hidden_states,+                )+            )+         if streamer is not None:             streamer.end() 

@@ -71,6 +71,10 @@ def extra_repr(self):   class MiniMaxCache(DynamicCache):+    # `crop` raises below, so a rollback cannot be undone here. The inherited property only inspects+    # `self.layers`, which are all plain dynamic layers, and would otherwise wrongly report `True`.+    is_croppable = False+     def __init__(self):         super().__init__()         self.linear_cache: list[torch.Tensor] = []

@@ -161,6 +161,10 @@ class MiniMaxRMSNorm(MixtralRMSNorm):   class MiniMaxCache(DynamicCache):+    # `crop` raises below, so a rollback cannot be undone here. The inherited property only inspects+    # `self.layers`, which are all plain dynamic layers, and would otherwise wrongly report `True`.+    is_croppable = False+     def __init__(self):         super().__init__()         self.linear_cache: list[torch.Tensor] = []

@@ -114,7 +114,7 @@         AssistedCandidateGeneratorDifferentTokenizers,         DFlashTokenCandidateGenerator,     )-    from transformers.generation.utils import ALL_CACHE_NAMES, _speculative_sampling+    from transformers.generation.utils import ALL_CACHE_NAMES, DeferredStopCheck, _speculative_sampling     from transformers.modeling_layers import MtpModel  from unittest.mock import patch@@ -5700,3 +5700,105 @@ def assert_similar_generate_outputs(output_1, output_2, atol=1e-5, rtol=1e-5):     mismatch_message = _get_generate_outputs_mismatch_message(output_1, output_2, atol=atol, rtol=rtol)     if mismatch_message:         raise AssertionError(mismatch_message)+++class _CollectingStreamer:+    """Keeps every tensor it is handed, so a stream can be compared against the sequences that were returned."""++    def __init__(self):+        self.tokens = []++    def put(self, value):+        self.tokens.append(value.reshape(-1).cpu())++    def end(self):+        pass+++@require_torch_accelerator+class DeferredStopCheckIntegrationTest(unittest.TestCase):+    """Deferring the stop decision must not change a single token of what `generate` returns."""++    model_id = "hf-internal-testing/tiny-random-gpt2"++    def setUp(self):+        self.tokenizer = AutoTokenizer.from_pretrained(self.model_id)+        self.model = AutoModelForCausalLM.from_pretrained(self.model_id).to(torch_device).eval()++    def _generate(self, inputs, model=None, **kwargs):+        """Generate twice over the same inputs, once collecting the stream, and return both."""+        model = self.model if model is None else model+        kwargs = {"max_new_tokens": 12, "do_sample": False, **kwargs}+        outputs = model.generate(**inputs, return_dict_in_generate=True, output_scores=True, **kwargs)+        streamer = _CollectingStreamer()+        model.generate(**inputs, streamer=streamer, **kwargs)+        return outputs, torch.cat(streamer.tokens)++    def _generate_deferred_and_immediate(self, inputs, model=None, **kwargs):+        # Both halves are forced rather than left to `is_supported`, so this holds on any accelerator:+        # `is_supported` limits where the deferral is *enabled*, but the logic under test is not+        # device-specific and is worth checking wherever the suite runs.+        with patch.object(DeferredStopCheck, "is_supported", staticmethod(lambda *args, **kwargs: True)):+            deferred = self._generate(inputs, model=model, **kwargs)+        with patch.object(DeferredStopCheck, "is_supported", staticmethod(lambda *args, **kwargs: False)):+            immediate = self._generate(inputs, model=model, **kwargs)+        return deferred, immediate++    def _assert_matches(self, inputs, model=None, **kwargs):+        prompt_length = inputs["input_ids"].shape[1]+        (deferred, stream), (immediate, _) = self._generate_deferred_and_immediate(inputs, model=model, **kwargs)++        # The extra step leaves no trace: same tokens, and one score per token actually returned+        self.assertTrue(torch.equal(deferred.sequences, immediate.sequences))+        self.assertEqual(len(deferred.scores), len(immediate.scores))+        self.assertEqual(len(deferred.scores), deferred.sequences.shape[1] - prompt_length)+        # A streamer is one step behind, but still sees every token and never the one past the stop+        self.assertEqual(stream.tolist(), deferred.sequences[0].cpu().tolist())++    def test_matches_immediate_check_at_max_length(self):+        inputs = self.tokenizer(["Hello world, this is"], return_tensors="pt").to(torch_device)+        self._assert_matches(inputs)++    def test_matches_immediate_check_when_stopping_early(self):+        """The step taken past an eos is the one that has to be undone, so this is the interesting case."""+        inputs = self.tokenizer(["Hello world, this is"], return_tensors="pt").to(torch_device)+        prompt_length = inputs["input_ids"].shape[1]+        # Make a token the model does generate into an eos, so generation really stops before `max_new_tokens`+        generated = self.model.generate(**inputs, max_new_tokens=12, do_sample=False)+        eos_token_id = int(generated[0, prompt_length + 4])++        self._assert_matches(inputs, eos_token_id=eos_token_id)+        stopped = self.model.generate(**inputs, max_new_tokens=12, do_sample=False, eos_token_id=eos_token_id)+        self.assertLess(stopped.shape[1], generated.shape[1])++    def test_matches_immediate_check_with_a_sliding_window(self):+        """A sliding window layer has to hold an extra state for the rollback without widening its window.++        `update` hands attention whatever the layer is holding, so a cache recording past for `crop` must+        still return the working window - otherwise the model silently attends over `sliding_window + 1`+        positions and generates different tokens.+        """+        model_id = "hf-internal-testing/tiny-random-Gemma3ForCausalLM"+        config = AutoConfig.from_pretrained(model_id)+        # Narrow the window so that a short generation still runs past it, which is the only regime where the+        # states are trimmed, and so the only one where holding an extra one could widen what attention sees+        config.sliding_window = 4+        model = AutoModelForCausalLM.from_pretrained(model_id, config=config).to(torch_device).eval()+        tokenizer = AutoTokenizer.from_pretrained(model_id)+        inputs = tokenizer(["Hello world, this is"], return_tensors="pt").to(torch_device)++        self._assert_matches(inputs, model=model)++    def test_matches_immediate_check_without_a_cache(self):+        inputs = self.tokenizer(["Hello world, this is"], return_tensors="pt").to(torch_device)+        self._assert_matches(inputs, use_cache=False)++    def test_matches_immediate_check_on_a_batch(self):+        self.tokenizer.pad_token = self.tokenizer.eos_token+        inputs = self.tokenizer(+            ["Hello world, this is", "The capital of France is"], return_tensors="pt", padding=True+        ).to(torch_device)+        prompt_length = inputs["input_ids"].shape[1]+        (deferred, _), (immediate, _) = self._generate_deferred_and_immediate(inputs)+        self.assertTrue(torch.equal(deferred.sequences, immediate.sequences))+        self.assertEqual(len(deferred.scores), deferred.sequences.shape[1] - prompt_length)