logit processor not behaving as expected

[jwd-dev]

I have a really simple logit processor implementation here but it's not behaving the way I would expect and I'm having trouble finding docs on how it's supposed to behave.

import torch
from typing import List, Optional

class DebugLogitProcessor:
    def __init__(self, tokenizer):
        self.tokenizer = tokenizer

    def __call__(self, req_id: int, logits: torch.Tensor, token_ids: List[List[int]], stream_ptr: Optional[int] = None, client_id: Optional[int] = None):
        print(f"[DebugLogitProcessor] req_id: {req_id}")
        logits_flat = logits[0, 0]
        top_idx = torch.argmax(logits_flat).item()
        token_str = self.tokenizer.decode([top_idx])
        print(f"[DebugLogitProcessor] Top token id={top_idx} token='{token_str}' logit={logits_flat[top_idx].item():.5f}")
        if token_ids:
            for i, seq in enumerate(token_ids):
                last_50 = seq[-50:] if len(seq) > 50 else seq
                decoded = self.tokenizer.decode(last_50)
                print(f"[DebugLogitProcessor] Beam {i}: Last 50 decoded: '{decoded}'")

        print(f"[DebugLogitProcessor] stream_ptr: {stream_ptr}")
        print(f"[DebugLogitProcessor] client_id: {client_id}")

Here's my logit processor I'm passing to the SamplingParams which then goes to generate_async. I'm using a pytorch backend and the executor API is enabled.

When generating here's the sequence I'm seeing:

This does not behave how I expected, I thought the logit processor is called once per forward pass and whatever the "chosen" token was after the forward pass would show up in the token_ids array during the next forward pass. (I'm sorry if this contains any misunderstandings, would love to be pointed in the right direction)

This makes it hard to implement something like: If the model just output the token "blue" force the next token to be "berry" As there's no easy way to know the current state of the generation.

Anyone have thoughts or pointers?