Keep image inputs on the first chunk of a chunked prefill

@@ -174,7 +174,6 @@     "patches_grid",     "image_token_pooling",     "image_grids",-    "image_num_crops",     "video_token_pooling",     "video_grids", )@@ -701,11 +700,12 @@ def prepare_inputs_for_generation(         position_ids_key = "decoder_position_ids" if self.config.is_encoder_decoder else "position_ids"         if (position_ids := kwargs.pop(position_ids_key, None)) is not None:             model_inputs[position_ids_key] = position_ids-        if (token_type_ids := kwargs.pop("token_type_ids", None)) is not None:-            model_inputs["token_type_ids"] = token_type_ids+        for token_type_key in ("token_type_ids", "mm_token_type_ids", "moe_mm_token_type_ids"):+            if (token_type := kwargs.pop(token_type_key, None)) is not None:+                model_inputs[token_type_key] = token_type          # 3. Slice model inputs if it's an input that should have the same length as `input_ids`-        for model_input_name in [position_ids_key, "token_type_ids", "mm_token_type_ids"]:+        for model_input_name in [position_ids_key, "token_type_ids", "mm_token_type_ids", "moe_mm_token_type_ids"]:             model_input = model_inputs.get(model_input_name)             if model_input is not None and model_input.shape[-1] != sequence_length:                 # Input can be 2D or 3D, and we always slice on `seq-length` (last dim)@@ -4312,7 +4312,18 @@ def _prefill(             chunk_size = generation_config.prefill_chunk_size             # In case we are restarting from non-empty Cache and the input_ids also contain the old ones, we need to slice             next_sequence_length = 0 if next_sequence_length is None else next_sequence_length-            input_chunks = torch.split(input_ids[:, -next_sequence_length:], chunk_size, dim=-1)+            prefill_ids = input_ids[:, -next_sequence_length:]+            input_chunks = torch.split(prefill_ids, chunk_size, dim=-1)+            raw_mm_keys = [+                key+                for key in MULTIMODAL_INPUTS_TO_DROP_OUTSIDE_PREFILL+                if key != "mm_encoder_outputs" and model_kwargs.get(key) is not None+            ]+            if is_first_iteration and len(input_chunks) > 1 and raw_mm_keys:+                raise ValueError(+                    f"Chunked prefill cannot split the raw multimodal inputs {raw_mm_keys} into chunks, so pass "+                    "`prefill_chunk_size=None` or use a model that pre-encodes them."+                )              model_forward = (                 self.get_compiled_call(generation_config.compile_config)@@ -4322,15 +4333,47 @@ def _prefill(              attention_mask = model_kwargs.pop("attention_mask", None)             position_ids = model_kwargs.pop("position_ids", None)+            token_types = {+                key: model_kwargs.pop(key)+                for key in ("token_type_ids", "mm_token_type_ids", "moe_mm_token_type_ids")+                if model_kwargs.get(key) is not None+            }+            mm_encoder_outputs = model_kwargs.pop("mm_encoder_outputs", None)++            # Pre-encoded features are consumed in placeholder order over (batch, seq): feature row `i` of a modality+            # fills its `i`-th placeholder token, so each chunk gets the rows of the placeholders it contains+            chunk_mm_encoder_outputs = [{} for _ in input_chunks]+            if is_first_iteration and mm_encoder_outputs:+                for modality, output in mm_encoder_outputs.items():+                    if output is None:+                        continue+                    mask = prefill_ids == getattr(self.config, f"{modality}_token_id")+                    feature_rows = mask.flatten().cumsum(0).view_as(mask) - 1+                    for chunk_idx, (chunk_mask, chunk_rows) in enumerate(+                        zip(mask.split(chunk_size, dim=-1), feature_rows.split(chunk_size, dim=-1))+                    ):+                        if chunk_mask.any():+                            chunk_mm_encoder_outputs[chunk_idx][modality] = _slice_mm_encoder_output(+                                output, chunk_rows[chunk_mask]+                            )+             past_length = cache.get_seq_length()-            for input_chunk in input_chunks:+            for chunk_idx, input_chunk in enumerate(input_chunks):                 current_length = past_length + input_chunk.shape[-1]                 if attention_mask is not None:                     model_kwargs["attention_mask"] = attention_mask[:, :current_length]                 if position_ids is not None:                     # Always slice on `seq-length` (last dim)                     model_kwargs["position_ids"] = position_ids[..., past_length:current_length]-                model_inputs = self.prepare_inputs_for_generation(input_chunk, **model_kwargs)+                # Up to the chunk end like the attention mask, for the models that read them by position+                for key, token_type in token_types.items():+                    model_kwargs[key] = token_type[..., :current_length]+                model_inputs = self.prepare_inputs_for_generation(+                    input_chunk, is_first_iteration=is_first_iteration, **model_kwargs+                )++                if chunk_mm_encoder_outputs[chunk_idx]:+                    model_inputs["mm_encoder_outputs"] = chunk_mm_encoder_outputs[chunk_idx]                  outputs = model_forward(**model_inputs, return_dict=True) @@ -4340,11 +4383,36 @@ def _prefill(             # Recreate the kwargs based on the full length             model_kwargs["attention_mask"] = attention_mask             model_kwargs["position_ids"] = position_ids+            model_kwargs.update(token_types)+            if mm_encoder_outputs is not None:+                model_kwargs["mm_encoder_outputs"] = mm_encoder_outputs              # Latest outputs contain next token logits             return outputs  +def _slice_mm_encoder_output(output: ModelOutput, rows: torch.LongTensor) -> ModelOutput:+    """+    Keeps the given rows of the flattened `(num_placeholder_tokens, hidden_size)` features of a pre-encoded multimodal+    output, in the container type the model expects: a single-item list, or a tensor with a batch of one. Only the+    pooled and deepstack features are kept, as those are what the models read from pre-encoded outputs.+    """++    def slice_features(features):+        rows_features = torch.cat([feature.reshape(-1, feature.shape[-1]) for feature in features])[rows]+        if isinstance(features, torch.Tensor):+            return rows_features.unsqueeze(0)+        return [rows_features]++    sliced = {"pooler_output": slice_features(output["pooler_output"])}+    deepstack_features = output.get("deepstack_features")+    if deepstack_features is not None:+        if isinstance(deepstack_features[0][0], torch.Tensor):+            deepstack_features = [slice_features(layer_features) for layer_features in deepstack_features]+        sliced["deepstack_features"] = deepstack_features+    return type(output)(**sliced)++ def _speculative_sampling(     candidate_input_ids,     candidate_logits,

@@ -1034,7 +1034,10 @@ def prepare_inputs_for_generation(self, input_ids, use_cache=True, is_first_iter         model_inputs = super().prepare_inputs_for_generation(             input_ids, use_cache=use_cache, is_first_iteration=is_first_iteration, **kwargs         )-        if not (is_first_iteration or not use_cache):+        if is_first_iteration or not use_cache:+            # The bidirectional mask reads them by position, so a prefill chunk needs them from the start of the sequence+            model_inputs["token_type_ids"] = kwargs.get("token_type_ids")+        else:             # Don't pass to not apply bidirectional mask on top             model_inputs["token_type_ids"] = None 

@@ -858,7 +858,10 @@ def prepare_inputs_for_generation(self, input_ids, use_cache=True, is_first_iter         model_inputs = super().prepare_inputs_for_generation(             input_ids, use_cache=use_cache, is_first_iteration=is_first_iteration, **kwargs         )-        if not (is_first_iteration or not use_cache):+        if is_first_iteration or not use_cache:+            # The bidirectional mask reads them by position, so a prefill chunk needs them from the start of the sequence+            model_inputs["token_type_ids"] = kwargs.get("token_type_ids")+        else:             # Don't pass to not apply bidirectional mask on top             model_inputs["token_type_ids"] = None 

@@ -2677,7 +2677,10 @@ def prepare_inputs_for_generation(self, input_ids, use_cache=True, is_first_iter         model_inputs = super().prepare_inputs_for_generation(             input_ids, use_cache=use_cache, is_first_iteration=is_first_iteration, **kwargs         )-        if not (is_first_iteration or not use_cache):+        if is_first_iteration or not use_cache:+            # The bidirectional mask reads them by position, so a prefill chunk needs them from the start of the sequence+            model_inputs["mm_token_type_ids"] = kwargs.get("mm_token_type_ids")+        else:             # Don't pass to not apply bidirectional mask on top             model_inputs["mm_token_type_ids"] = None         # If `per_layer_inputs` was provided along with `inputs_embeds` for first forward, drop it for subsequent forwards

@@ -2349,7 +2349,10 @@ def prepare_inputs_for_generation(self, input_ids, use_cache=True, is_first_iter         model_inputs = super().prepare_inputs_for_generation(             input_ids, use_cache=use_cache, is_first_iteration=is_first_iteration, **kwargs         )-        if not (is_first_iteration or not use_cache):+        if is_first_iteration or not use_cache:+            # The bidirectional mask reads them by position, so a prefill chunk needs them from the start of the sequence+            model_inputs["mm_token_type_ids"] = kwargs.get("mm_token_type_ids")+        else:             # Don't pass to not apply bidirectional mask on top             model_inputs["mm_token_type_ids"] = None         # If `per_layer_inputs` was provided along with `inputs_embeds` for first forward, drop it for subsequent forwards

@@ -1414,7 +1414,10 @@ def prepare_inputs_for_generation(self, input_ids, use_cache=True, is_first_iter         model_inputs = super().prepare_inputs_for_generation(             input_ids, use_cache=use_cache, is_first_iteration=is_first_iteration, **kwargs         )-        if not (is_first_iteration or not use_cache):+        if is_first_iteration or not use_cache:+            # The bidirectional mask reads them by position, so a prefill chunk needs them from the start of the sequence+            model_inputs["mm_token_type_ids"] = kwargs.get("mm_token_type_ids")+        else:             # Don't pass to not apply bidirectional mask on top             model_inputs["mm_token_type_ids"] = None 

@@ -1312,7 +1312,10 @@ def prepare_inputs_for_generation(self, input_ids, use_cache=True, is_first_iter         model_inputs = super().prepare_inputs_for_generation(             input_ids, use_cache=use_cache, is_first_iteration=is_first_iteration, **kwargs         )-        if not (is_first_iteration or not use_cache):+        if is_first_iteration or not use_cache:+            # The bidirectional mask reads them by position, so a prefill chunk needs them from the start of the sequence+            model_inputs["mm_token_type_ids"] = kwargs.get("mm_token_type_ids")+        else:             # Don't pass to not apply bidirectional mask on top             model_inputs["mm_token_type_ids"] = None 

@@ -1035,6 +1035,23 @@ def forward(             image_hidden_states=outputs.image_hidden_states,         ) +    def prepare_inputs_for_generation(+        self, input_ids, inputs_embeds=None, use_cache=True, is_first_iteration=False, **kwargs+    ):+        model_inputs = super().prepare_inputs_for_generation(+            input_ids,+            inputs_embeds=inputs_embeds,+            use_cache=use_cache,+            is_first_iteration=is_first_iteration,+            **kwargs,+        )+        if is_first_iteration or not use_cache:+            # The block mask reads them by position, so a prefill chunk needs them from the start of the sequence+            model_inputs["mm_token_type_ids"] = kwargs.get("mm_token_type_ids")+        else:+            model_inputs["mm_token_type_ids"] = None+        return model_inputs+     @staticmethod     def create_masks_for_generate(         config: PreTrainedConfig,

@@ -1781,6 +1781,23 @@ def forward(             image_hidden_states=outputs.image_hidden_states,         ) +    def prepare_inputs_for_generation(+        self, input_ids, inputs_embeds=None, use_cache=True, is_first_iteration=False, **kwargs+    ):+        model_inputs = super().prepare_inputs_for_generation(+            input_ids,+            inputs_embeds=inputs_embeds,+            use_cache=use_cache,+            is_first_iteration=is_first_iteration,+            **kwargs,+        )+        if is_first_iteration or not use_cache:+            # The block mask reads them by position, so a prefill chunk needs them from the start of the sequence+            model_inputs["mm_token_type_ids"] = kwargs.get("mm_token_type_ids")+        else:+            model_inputs["mm_token_type_ids"] = None+        return model_inputs+     @staticmethod     def create_masks_for_generate(         config: PreTrainedConfig,

@@ -48,19 +48,33 @@         AutoModelForCausalLM,         AutoTokenizer,         Cache,+        CLIPVisionConfig,+        CTRLConfig,+        CTRLLMHeadModel,         DynamicCache,+        Ernie4_5_VLMoeConfig,+        Ernie4_5_VLMoeForConditionalGeneration,         Gemma2Config,+        Gemma3Config,+        Gemma3ForConditionalGeneration,+        Gemma3TextConfig,+        Gemma4Config,         Gemma4ForCausalLM,+        Gemma4ForConditionalGeneration,         Gemma4TextConfig,+        Gemma4VisionConfig,         GenerationConfig,         LlamaConfig,         LlamaForCausalLM,+        LlavaConfig,+        LlavaForConditionalGeneration,         MixtralConfig,         MixtralForCausalLM,         PreTrainedConfig,         QuantizedCache,         Qwen3VLConfig,         Qwen3VLForConditionalGeneration,+        SiglipVisionConfig,         StaticCache,         convert_and_export_with_cache,         pipeline,@@ -519,6 +533,288 @@ def test_chunked_prefill_cache_continuation(self, cache_implementation):         self.assertTrue(torch.equal(no_chunked_only_new_ids, chunked_only_new_ids))         self.assertTrue(torch.equal(chunked_only_new_ids, chunked_full_ids[:, -chunked_only_new_ids.shape[-1] :])) +    @staticmethod+    def _get_tiny_llava():+        set_seed(3)+        config = LlavaConfig(+            vision_config=CLIPVisionConfig(+                hidden_size=16,+                intermediate_size=32,+                num_hidden_layers=1,+                num_attention_heads=2,+                image_size=8,+                patch_size=4,+                projection_dim=16,+            ),+            text_config=LlamaConfig(+                vocab_size=64,+                hidden_size=16,+                intermediate_size=32,+                num_hidden_layers=1,+                num_attention_heads=2,+                num_key_value_heads=2,+                max_position_embeddings=64,+                pad_token_id=0,+            ),+            image_token_index=63,+            vision_feature_layer=-1,+            vision_feature_select_strategy="full",+        )+        return LlavaForConditionalGeneration(config).to(torch_device).eval()++    def test_chunked_prefill_keeps_image_inputs(self):+        """+        A chunked prefill must pass the image inputs to its first forward, as the unchunked prefill does. They were+        dropped from every chunk, so a vision-language model generated without the image.+        """+        model = self._get_tiny_llava()+        # 5 image tokens (4 patches + CLS) in the first chunk of 6, then 4 text tokens+        input_ids = torch.tensor([[1] + [63] * 5 + [5, 6, 7, 8]], device=torch_device)+        image_a = torch.randn(1, 3, 8, 8, device=torch_device)+        image_b = torch.randn(1, 3, 8, 8, device=torch_device) * 5+        generation_kwargs = {+            "max_new_tokens": 6,+            "do_sample": False,+            "output_scores": True,+            "return_dict_in_generate": True,+        }++        expected = model.generate(input_ids, pixel_values=image_a, **generation_kwargs)+        out = model.generate(input_ids, pixel_values=image_a, **generation_kwargs, prefill_chunk_size=6)+        out_b = model.generate(input_ids, pixel_values=image_b, **generation_kwargs, prefill_chunk_size=6)++        self.assertTrue(torch.equal(out.sequences, expected.sequences))+        torch.testing.assert_close(torch.stack(out.scores), torch.stack(expected.scores))+        self.assertFalse(torch.equal(torch.stack(out.scores), torch.stack(out_b.scores)))++    def test_chunked_prefill_splits_image_features(self):+        """+        A chunk boundary may fall inside a run of image tokens: each chunk must get the image features of its own image+        tokens, so that any `prefill_chunk_size` gives the same output as the unchunked prefill.+        """+        model = self._get_tiny_llava()+        generation_kwargs = {+            "max_new_tokens": 6,+            "do_sample": False,+            "output_scores": True,+            "return_dict_in_generate": True,+        }+        one_image = torch.tensor([[1] + [63] * 5 + [5, 6, 7, 8]], device=torch_device)+        # Image runs at positions 1-5 and 8-12: chunks of 4 split both of them+        two_images = torch.tensor([[1] + [63] * 5 + [5, 6] + [63] * 5 + [7, 8]], device=torch_device)+        two_prompts = torch.tensor(+            [[1] + [63] * 5 + [5, 6, 7, 8], [1, 9, 10] + [63] * 5 + [11, 12]], device=torch_device+        )+        cases = [+            (one_image, torch.randn(1, 3, 8, 8, device=torch_device), [4, 2]),+            (two_images, torch.randn(2, 3, 8, 8, device=torch_device), [4, 3]),+            (two_prompts, torch.randn(2, 3, 8, 8, device=torch_device), [4]),+        ]+        for input_ids, pixel_values, chunk_sizes in cases:+            expected = model.generate(input_ids, pixel_values=pixel_values, **generation_kwargs)+            for chunk_size in chunk_sizes:+                case = f"{pixel_values.shape[0]} images, batch {input_ids.shape[0]}, prefill_chunk_size={chunk_size}"+                out = model.generate(+                    input_ids, pixel_values=pixel_values, **generation_kwargs, prefill_chunk_size=chunk_size+                )+                self.assertTrue(torch.equal(out.sequences, expected.sequences), case)+                torch.testing.assert_close(torch.stack(out.scores), torch.stack(expected.scores), msg=case)++    def test_chunked_prefill_raises_on_raw_image_inputs(self):+        """+        Raw image inputs cannot be split into chunks: a model that does not pre-encode them must raise when the prefill+        is split, and still work when the prefill fits in one chunk.+        """+        model = self._get_tiny_llava()+        model._supports_mm_encoder_outputs = lambda: False+        input_ids = torch.tensor([[1] + [63] * 5 + [5, 6, 7, 8]], device=torch_device)+        pixel_values = torch.randn(1, 3, 8, 8, device=torch_device)+        generation_kwargs = {+            "max_new_tokens": 6,+            "do_sample": False,+            "output_scores": True,+            "return_dict_in_generate": True,+        }++        with self.assertRaisesRegex(ValueError, "pixel_values"):+            model.generate(input_ids, pixel_values=pixel_values, **generation_kwargs, prefill_chunk_size=6)++        expected = model.generate(input_ids, pixel_values=pixel_values, **generation_kwargs)+        out = model.generate(+            input_ids, pixel_values=pixel_values, **generation_kwargs, prefill_chunk_size=input_ids.shape[1]+        )+        self.assertTrue(torch.equal(out.sequences, expected.sequences))+        torch.testing.assert_close(torch.stack(out.scores), torch.stack(expected.scores))++    def test_chunked_prefill_keeps_token_type_ids(self):+        """+        Gemma3 and Gemma4 build a bidirectional mask over each image run from `token_type_ids` and `mm_token_type_ids`.+        Every chunk of a chunked prefill must get them, or the image tokens of a later chunk attend causally only and+        the output differs from the unchunked prefill.+        """+        set_seed(3)+        config = Gemma3Config(+            vision_config=SiglipVisionConfig(+                hidden_size=16,+                intermediate_size=32,+                num_hidden_layers=1,+                num_attention_heads=2,+                image_size=8,+                patch_size=2,+            ),+            text_config=Gemma3TextConfig(+                vocab_size=64,+                hidden_size=16,+                intermediate_size=32,+                num_hidden_layers=2,+                num_attention_heads=2,+                num_key_value_heads=2,+                head_dim=8,+                max_position_embeddings=64,+                sliding_window=16,+                layer_types=["sliding_attention", "full_attention"],+                pad_token_id=0,+            ),+            mm_tokens_per_image=4,+            boi_token_index=61,+            eoi_token_index=62,+            image_token_index=63,+        )+        model = Gemma3ForConditionalGeneration(config).to(torch_device).eval()+        generation_kwargs = {+            "max_new_tokens": 6,+            "do_sample": False,+            "output_scores": True,+            "return_dict_in_generate": True,+        }+        image = [61] + [63] * 4 + [62]+        # Chunks of 4: the image run fills the second chunk+        later_image = torch.tensor([[1, 2, 3] + image + [5, 6, 7]], device=torch_device)+        # Chunks of 8: one image run in each chunk+        two_images = torch.tensor([[1] + image + [5, 6] + image + [7]], device=torch_device)+        cases = [+            (later_image, torch.randn(1, 3, 8, 8, device=torch_device), 4),+            (two_images, torch.randn(2, 3, 8, 8, device=torch_device), 8),+        ]+        for input_ids, pixel_values, chunk_size in cases:+            inputs = {"pixel_values": pixel_values, "token_type_ids": (input_ids == 63).long()}+            expected = model.generate(input_ids, **inputs, **generation_kwargs)+            out = model.generate(input_ids, **inputs, **generation_kwargs, prefill_chunk_size=chunk_size)+            case = f"{pixel_values.shape[0]} images, prefill_chunk_size={chunk_size}"+            self.assertTrue(torch.equal(out.sequences, expected.sequences), case)+            torch.testing.assert_close(torch.stack(out.scores), torch.stack(expected.scores), msg=case)++        config = Gemma4Config(+            vision_config=Gemma4VisionConfig(+                hidden_size=16,+                intermediate_size=32,+                num_hidden_layers=1,+                num_attention_heads=2,+                num_key_value_heads=2,+                head_dim=8,+                patch_size=2,+                pooling_kernel_size=2,+                position_embedding_size=16,+            ),+            text_config=Gemma4TextConfig(+                vocab_size=64,+                hidden_size=16,+                intermediate_size=32,+                num_hidden_layers=2,+                num_attention_heads=2,+                num_key_value_heads=2,+                head_dim=8,+                max_position_embeddings=64,+                sliding_window=16,+                layer_types=["sliding_attention", "full_attention"],+                vocab_size_per_layer_input=64,+                hidden_size_per_layer_input=8,+                use_bidirectional_attention="vision",+                pad_token_id=0,+            ),+            boi_token_id=61,+            eoi_token_id=62,+            image_token_id=63,+            video_token_id=60,+            audio_token_id=59,+        )+        model = Gemma4ForConditionalGeneration(config).to(torch_device).eval()+        # A 4x4 grid of patches, pooled 2x2 into the 4 image tokens of the second chunk+        grid = torch.stack([torch.arange(4).repeat(4), torch.arange(4).repeat_interleave(4)], dim=-1)+        inputs = {+            "pixel_values": torch.randn(1, 16, 12, device=torch_device),+            "image_position_ids": grid[None].to(torch_device),+            "mm_token_type_ids": (later_image == 63).long(),+        }+        expected = model.generate(later_image, **inputs, **generation_kwargs)+        out = model.generate(later_image, **inputs, **generation_kwargs, prefill_chunk_size=4)+        self.assertTrue(torch.equal(out.sequences, expected.sequences))+        torch.testing.assert_close(torch.stack(out.scores), torch.stack(expected.scores))++    def test_chunked_prefill_slices_per_token_types(self):+        """+        Ernie4.5-VL-MoE routes each token by its `moe_mm_token_type_ids`, so every chunk of a chunked prefill must+        get the types of its own tokens. CTRL drops `token_type_ids` in generation, so no chunk may get them back.+        """+        set_seed(3)+        # Both layers route by type, so a wrong route of the first one reaches the later chunks through the cache+        config = Ernie4_5_VLMoeConfig(+            text_config={+                "vocab_size": 64,+                "hidden_size": 32,+                "intermediate_size": 32,+                "num_hidden_layers": 2,+                "num_attention_heads": 2,+                "num_key_value_heads": 1,+                "rope_parameters": {"rope_type": "default", "rope_theta": 500_000.0, "mrope_section": [3, 3, 2]},+                "moe_intermediate_size": [32, 32],+                "mlp_layer_types": ["sparse", "sparse"],+                "pad_token_id": 0,+            },+            vision_config={+                "depth": 1,+                "hidden_size": 32,+                "intermediate_size": 32,+                "num_heads": 2,+                "spatial_merge_size": 1,+            },+            image_start_token_id=5,+            image_end_token_id=6,+            image_token_id=7,+            video_start_token_id=3,+            video_end_token_id=4,+            video_token_id=8,+        )+        model = Ernie4_5_VLMoeForConditionalGeneration(config).to(torch_device).eval()+        generation_kwargs = {+            "max_new_tokens": 6,+            "do_sample": False,+            "output_scores": True,+            "return_dict_in_generate": True,+        }+        # Chunks of 4: the 2x2 patches of the image fill the second chunk+        input_ids = torch.tensor([[1, 9, 10, 5] + [7] * 4 + [6, 11, 12, 13]], device=torch_device)+        inputs = {+            "pixel_values": torch.randn(4, 3 * 14 * 14, device=torch_device),+            "image_grid_thw": torch.tensor([[1, 2, 2]], device=torch_device),+            "mm_token_type_ids": (input_ids == 7).long(),+            "moe_mm_token_type_ids": torch.isin(input_ids, torch.tensor([5, 6, 7], device=torch_device)).long(),+        }+        expected = model.generate(input_ids, **inputs, **generation_kwargs)+        out = model.generate(input_ids, **inputs, **generation_kwargs, prefill_chunk_size=4)+        self.assertTrue(torch.equal(out.sequences, expected.sequences))+        torch.testing.assert_close(torch.stack(out.scores), torch.stack(expected.scores))++        set_seed(3)+        config = CTRLConfig(vocab_size=64, n_embd=16, dff=32, n_layer=1, n_head=2, n_positions=64, pad_token_id=0)+        model = CTRLLMHeadModel(config).to(torch_device).eval()+        input_ids = torch.tensor([[1, 9, 10, 11, 12, 13, 14, 15, 16, 17]], device=torch_device)+        token_type_ids = torch.tensor([[0, 0, 1, 1, 0, 0, 1, 1, 0, 0]], device=torch_device)+        expected = model.generate(input_ids, token_type_ids=token_type_ids, **generation_kwargs)+        out = model.generate(input_ids, token_type_ids=token_type_ids, **generation_kwargs, prefill_chunk_size=4)+        self.assertTrue(torch.equal(out.sequences, expected.sequences))+        torch.testing.assert_close(torch.stack(out.scores), torch.stack(expected.scores))+     def test_dynamic_layers_reset_drops_their_states(self):         """         Regression test: `reset` used to zero the dynamic layers' states in place, inherited from the static layers.