Re-enable prefill of max model length

[yannicks1]

Re-enable prefill of max model length

Purpose

Previous to #20291 it was possible to do a prefill on the max model length and sample a single token. However, #20291 introduced an assertion:

start_idx = self.input_batch.num_tokens_no_spec[req_idx]
end_idx = start_idx + len(sampled_ids)
assert end_idx <= self.max_model_len <-----

which does not allow this anymore.

This PR allows the initial behavior again. Note that this initial behavior (allowing to do a prefill and requesting a single token on the max model length) is consistent with huggingface transformer behavior.

Test Plan

Can be tested with any decoder model with a prompt of length max_model_len and setting max_tokens=1.

Test Result

e.g before this PR max_model_len=2048 and max_tokens=1 resulted in this warning:

AssertionError: Sampled token IDs exceed the max model length. Total number of tokens: 2049 > max_model_len: 2048

after this PR: the assertion error does not show up and the correct output token is returned.

[gemini-code-assist]

Code Review

This pull request relaxes an assertion to re-enable prefilling up to the maximum model length and sampling a single token. While the intent is correct, the change as-is will likely cause an IndexError because the underlying buffer for token IDs is not large enough to accommodate the extra token. A fix is required in vllm/v1/worker/gpu_input_batch.py (and likely vllm/v1/worker/tpu_input_batch.py) to increase the buffer size.

[gemini-code-assist]

While this change correctly relaxes the assertion to allow for prefilling up to max_model_len and sampling one more token, it seems to introduce a potential IndexError in the subsequent line.

The buffer self.input_batch.token_ids_cpu is initialized with a size of (max_num_reqs, max_model_len) in vllm/v1/worker/gpu_input_batch.py.

When end_idx is self.max_model_len + 1, the slice assignment self.input_batch.token_ids_cpu[req_idx, start_idx:end_idx] = sampled_ids will attempt to write to an out-of-bounds index (max_model_len).

To fix this, the token_ids_cpu buffer should probably be initialized with a size of (max_num_reqs, max_model_len + 1). This change would be required in vllm/v1/worker/gpu_input_batch.py. A similar change might be needed for vllm/v1/worker/tpu_input_batch.py as well.

[yannicks1]

Not sure if the change in vllm/v1/worker/tpu_input_batch.py is needed.

[yannicks1]

@WoosukKwon @LucasWilkinson tagging you guys here as author/reviewer of #20291