[BugFix] Pad input buffers in _dummy_run

[varun-sundar-rabindranath]

Purpose

Fixes #26137

PR #24845 removed the padding. This PR re-introduce the padding.

Test Plan

CI & commands in the issue run successfully.

Test Result

[gemini-code-assist]

Code Review

This pull request re-introduces padding for dummy runs, which is a necessary bug fix. The changes correctly replace num_tokens with num_tokens_after_padding in various places. However, I've identified a potential critical issue where num_tokens_after_padding can exceed the allocated buffer size, leading to an out-of-bounds memory access.

[gemini-code-assist]

There is a potential out-of-bounds memory access issue here and in the following lines that use num_tokens_after_padding for slicing. The value of num_tokens_after_padding can exceed self.max_num_tokens, which is the size of buffers like self.input_ids, self.positions, and self.inputs_embeds.

Here's how it can happen in _dummy_run when DBO is enabled:

1. _dummy_run is called with num_tokens equal to self.max_num_tokens.
2. ubatch_split is called, which in turn calls get_dp_padding_ubatch.
3. get_dp_padding_ubatch calculates num_tokens_padded = round_up(num_tokens, 2). If self.max_num_tokens is odd, this results in self.max_num_tokens + 1.
4. This padded value is used to calculate num_tokens_per_ubatch, which is then communicated across DP ranks. The maximum is taken.
5. Back in _dummy_run, num_tokens_after_padding is calculated based on the result from ubatch_split, and can become self.max_num_tokens + 1.

Slicing tensors like self.input_ids.gpu[:num_tokens_after_padding] will then result in an out-of-bounds access, which can lead to memory corruption or a crash. This is a critical issue that needs to be addressed.

[varun-sundar-rabindranath]

REVIEWERS, PTAL

I have added an assert num_tokens_after_padding < self.max_num_tokens to address this.
A better fix is to roundup self.max_num_tokens here

vllm/vllm/v1/worker/gpu_model_runner.py

Line 222 in a42d2df

self.max_num_tokens = scheduler_config.max_num_batched_tokens

But we tend to check max_num_tokens against scheduler_config.max_num_batched_tokens in code and it is a reasonable check. This probably needs to be handled more carefully.

However, we generally don't expect max_num_tokens to be odd, and it might never be a issue. But when it does happen, this assert should catch it.

[varun-sundar-rabindranath]

Marking this as draft as this might not be the full fix and could introduce bugs in the DBO case.

cc @LucasWilkinson @SageMoore @ProExpertProg @ilmarkov @tlrmchlsmth

[SageMoore]

The biggest potential gotcha here is that num_tokens_after_padding will be divided in half for the DBO case. I suspect you will want to pad up the inputs to the "full" padded amount. The UbatchWrapper will take care of slicing them down to the ubatch sizes.

This is definitely annoying but once #25768 merges we won't have this difference.

[varun-sundar-rabindranath]

Marking this PR as draft (so we dont land it by mistake) as it deadlocks during DBO sanity checks (benchmarking)

[varun-sundar-rabindranath]

The biggest potential gotcha here is that num_tokens_after_padding will be divided in half for the DBO case. I suspect you will want to pad up the inputs to the "full" padded amount. The UbatchWrapper will take care of slicing them down to the ubatch sizes.

Thanks @SageMoore . I noticed the update to num_tokens_after_padding just before the model forward pass call. Should we do another slice of the buffers after that update? if we do that, I think it'll be correct for the DBO case and a no-op for the non DBO case.