@@ -318,8 +318,8 @@ class RepetitionPenaltyLogitsProcessor(LogitsProcessor):
         penalty (`float`):
             The parameter for repetition penalty. 1.0 means no penalty. Above 1.0 penalizes previously generated
             tokens. Between 0.0 and 1.0 rewards previously generated tokens.
-        prompt_ignore_length (`int`, *optional*):
-            The original input ids sequence length, which if provided, will not be used in the penalty calculation.
+        prompt_ignore_length (`int`, or `torch.LongTensor`, *optional*):
+            Number of leading tokens to ignore, optionally specified per batch element.
 
     Examples:
 
@@ -355,14 +355,15 @@ class RepetitionPenaltyLogitsProcessor(LogitsProcessor):
 
     supports_continuous_batching = False
 
-    def __init__(self, penalty: float, prompt_ignore_length: int | None = None):
+    def __init__(self, penalty: float, prompt_ignore_length: int | torch.Tensor | None = None):
         if not isinstance(penalty, float) or not (penalty > 0):
             raise ValueError(f"`penalty` has to be a strictly positive float, but is {penalty}")
 
-        if prompt_ignore_length is not None and (
-            not isinstance(prompt_ignore_length, int) or prompt_ignore_length < 0
-        ):
-            raise ValueError(f"`prompt_ignore_length` has to be a positive integer, but is {prompt_ignore_length}")
+        if prompt_ignore_length is not None:
+            if isinstance(prompt_ignore_length, int) and prompt_ignore_length < 0:
+                raise ValueError(f"`prompt_ignore_length` has to be a positive integer, but is {prompt_ignore_length}")
+            if isinstance(prompt_ignore_length, torch.Tensor) and prompt_ignore_length.dim() != 1:
+                raise ValueError("`prompt_ignore_length` must be 1D")
 
         self.penalty = penalty
         self.prompt_ignore_length = prompt_ignore_length
@@ -371,36 +372,49 @@ def __init__(self, penalty: float, prompt_ignore_length: int | None = None):
 
     @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING)
     def __call__(self, input_ids: torch.LongTensor, scores: torch.Tensor) -> torch.Tensor:
-        if self.prompt_ignore_length:
+        vocab_size = scores.shape[-1]
+        # In this case simply cut the prompt
+        if isinstance(self.prompt_ignore_length, int):
             input_ids = input_ids[:, self.prompt_ignore_length :]
+        # In this case, set all ids we do not want to take into account at `vocab_size`, we will then filter them out after `scatter`
+        elif isinstance(self.prompt_ignore_length, torch.Tensor):
+            if input_ids.shape[0] != self.prompt_ignore_length.shape[0]:
+                raise ValueError("Incompatible `prompt_ignore_length` batch size")
+            # Set the device in-place so that future calls will be a no-op
+            self.prompt_ignore_length = self.prompt_ignore_length.to(input_ids.device)
+            positions = torch.arange(input_ids.shape[1], device=input_ids.device)
+            input_ids = input_ids.masked_fill(positions[None, :] < self.prompt_ignore_length[:, None], vocab_size)
 
         if scores.dim() == 3:
             if self.logits_indices is not None and self.cu_seq_lens_q is not None:
                 last_positions = self.logits_indices
                 last_scores = scores[0, last_positions, :]
 
                 # Prepare token mask
-                token_mask = torch.zeros_like(last_scores, dtype=torch.bool)
+                token_mask = torch.zeros(
+                    (*last_scores.shape[:-1], vocab_size + 1), dtype=torch.bool, device=last_scores.device
+                )
                 cu_seq_lens = self.cu_seq_lens_q
                 lengths = cu_seq_lens[1:] - cu_seq_lens[:-1]
                 seq_indices = torch.repeat_interleave(torch.arange(len(lengths), device=input_ids.device), lengths)
                 token_mask[seq_indices, input_ids] = True
 
                 # Apply penalty
                 penalty_scores = torch.where(last_scores < 0, last_scores * self.penalty, last_scores / self.penalty)
-                scores[0, last_positions, :] = torch.where(token_mask, penalty_scores, last_scores)
+                scores[0, last_positions, :] = torch.where(token_mask[..., :vocab_size], penalty_scores, last_scores)
             else:
-                batch_size, seq_len, vocab_size = scores.shape
                 last_scores = scores[:, -1, :]
-                token_mask = torch.zeros_like(last_scores, dtype=torch.bool)
+                token_mask = torch.zeros(
+                    (*last_scores.shape[:-1], vocab_size + 1), dtype=torch.bool, device=last_scores.device
+                )
                 if input_ids.dim() == 1:
                     unique_tokens = torch.unique(input_ids)
                     token_mask.scatter_(1, unique_tokens.unsqueeze(0), True)
                 else:
                     token_mask.scatter_(1, input_ids, True)
                 # if last_scores < 0 then repetition penalty has to be multiplied to reduce the token probabilities
                 penalty_scores = torch.where(last_scores < 0, last_scores * self.penalty, last_scores / self.penalty)
-                scores[:, -1, :] = torch.where(token_mask, penalty_scores, last_scores)
+                scores[:, -1, :] = torch.where(token_mask[..., :vocab_size], penalty_scores, last_scores)
             return scores
 
         if input_ids.dim() == 1:
@@ -409,7 +423,6 @@ def __call__(self, input_ids: torch.LongTensor, scores: torch.Tensor) -> torch.T
         # In some very rare cases, the embedding vocab may be larger than the lm_head output - in those cases, clamp the
         # ids > output vocab_size to a fix index that we drop afterwards (the tokens with id > output_vocab cannot be
         # generated by the model, so there's no need to apply a penalty to them)
-        vocab_size = scores.shape[-1]
         input_ids = input_ids.clamp(max=vocab_size)
         # Pad by 1 to absorb all the indices larger than output vocab that were clamped to the vocab
         scores = torch.nn.functional.pad(scores, (0, 1))
 
@@ -1464,7 +1464,15 @@ def _get_logits_processor(
                     UserWarning,
                 )
         if generation_config.repetition_penalty is not None and generation_config.repetition_penalty != 1.0:
-            processors.append(RepetitionPenaltyLogitsProcessor(penalty=generation_config.repetition_penalty))
+            mask = (model_kwargs or {}).get(
+                "decoder_attention_mask" if self.config.is_encoder_decoder else "attention_mask"
+            )
+            padding_lengths = (mask != 0).int().argmax(dim=1) if mask is not None else None
+            processors.append(
+                RepetitionPenaltyLogitsProcessor(
+                    generation_config.repetition_penalty, prompt_ignore_length=padding_lengths
+                )
+            )
             if not self.config.is_encoder_decoder and (input_ids_seq_length is None or input_ids_seq_length == 0):
                 warnings.warn(
                     "Passing `repetition_penalty` with `inputs_embeds` and without `input_ids` to `generate` will "
 
@@ -23,6 +23,7 @@
     DiaEOSChannelFilterLogitsProcessor,
     DiaEOSDelayPatternLogitsProcessor,
     LogitsProcessorList,
+    RepetitionPenaltyLogitsProcessor,
     TemperatureLogitsWarper,
 )
 from ...generation.stopping_criteria import StoppingCriteriaList
@@ -85,6 +86,13 @@ def _get_logits_processor(
             negative_prompt_attention_mask=negative_prompt_attention_mask,
         )
 
+        # Dia expand inputs and thus cannot use tensor `prompt_ignore_length` correctly
+        for processor in merged_processors:
+            if isinstance(processor, RepetitionPenaltyLogitsProcessor) and isinstance(
+                getattr(processor, "prompt_ignore_length", None), torch.Tensor
+            ):
+                processor.prompt_ignore_length = None
+
         # Custom processors we need at specific positions
         if original_guidance_scale is not None and original_guidance_scale != 1:
             cfg_processor = DiaClassifierFreeGuidanceLogitsProcessor(
 
@@ -210,6 +210,19 @@ def test_repetition_penalty_dist_process(self):
         # processor should not change logits in-place
         self.assertFalse(torch.all(scores == processed_scores))
 
+    def test_repetition_penalty_with_batched_padding(self):
+        vocab_size = 100
+        ids = torch.randint(1, vocab_size - 1, (3, 15))
+        # Add padding of different length on each batch with ids == vocab_size for the one that are not padding, so they
+        # do not account in the scores
+        padding = torch.tensor([[0, 0, 0], [0, vocab_size, vocab_size], [0, 0, vocab_size]])
+        full_ids = torch.cat([padding, ids], dim=-1)
+        scores = torch.rand(3, vocab_size)
+        processor_with_length = RepetitionPenaltyLogitsProcessor(2.0, prompt_ignore_length=torch.tensor([3, 1, 2]))
+        processor_without_length = RepetitionPenaltyLogitsProcessor(2.0)
+        # Using the processor with ignore length on padded inputs should be the same as processor without on un-padded inputs
+        self.assertTrue((processor_with_length(full_ids, scores) == processor_without_length(ids, scores)).all())
+
     def test_repetition_penalty_dist_process_exclusion_no_new_input_ids(self):
         input_ids = torch.tensor([[0, 1], [5, 0]], device=torch_device, dtype=torch.long)
         vocab_size = 10