Vectorized (batched) logits processors for assisted decoding

Feature request

Vectorize the logits processor application over draft slots in assisted decoding, so speculative candidate generation can apply processors to all drafted positions in one call instead of looping one draft step at a time.

Context

This came up during review of #48281 (sampled DFlash returned a logits tensor missing the batch dimension). @zucchini-nlp asked about restructuring the candidate generation loop to apply logits processors in a single batched call, and noted a follow-up for this "would be really nice". @Cyrilvallez maintains candidate generation, so input on the design here would be very welcome before any code is written.

Today the assisted decoding loop re-feeds each sampled draft token and applies the processors one position at a time (see the per-slot loop around candidate_generator.py:1717). MtpModel.forward and PromptLookupCandidateGenerator have the same per-position structure.

What blocks a naive batched call

I audited all 35 concrete processors and warpers in logits_process.py for compatibility with (batch, draft_positions, vocab) input. Only 8 are safe to run on 3D input directly. The rest fall into two groups:

1. 13 crash or corrupt dimensions, for example:

   • TopPLogitsWarper hard-codes scatter(1, ...), so vocab ids index the time dim
   • EncoderRepetitionPenaltyLogitsProcessor gathers with a 2D index into 3D src
   • TopHLogitsWarper unpacks batch_size, vocab_size = scores.shape
   • WhisperTimeStampLogitsProcessor slices scores_processed[k, self.timestamp_begin:] across the time dim

2. 14 run but silently change semantics, because they read prefix state that differs per draft slot:

   • RepetitionPenaltyLogitsProcessor's existing 3D branch penalizes only scores[:, -1, :]
   • MinLengthLogitsProcessor / MinNewTokensLengthLogitsProcessor test input_ids.shape[-1] once and mask every position identically
   • NoRepeatNGramLogitsProcessor bans ngrams computed from the final suffix for all slots
   • ForcedEOSTokenLogitsProcessor fires at every position from a single length check

Each drafted slot needs its own growing prefix (the loop samples from one slot and appends before the next), so a single batch call cannot reproduce per-position semantics without making those processors position-aware.

Proposed direction

Two parts, kept separable:

1. Fix the dim-corrupting processors so they are correct on 3D input (scatter/gather dim inference, shape unpacks, time-dim slices). This is safe standalone and useful regardless.
2. Make prefix-state processors position-aware (accept per-position lengths or growing prefixes), and add a batched application path used by the candidate generator loop. The single-token path stays unchanged.

Who will potentially do this?

Happy to take this on. No code yet; waiting for @Cyrilvallez's and @zucchini-nlp's thoughts on scope and the preferred design before starting.