[Bugfix] Fix DP Attention Padding in Dummy Run

[LucasWilkinson]

Mirror of #34009 by @benchislett — "Maintainers are allowed to edit this pull request." was not enabled on the original PR, so pushing review fixes was not possible.

All credit to @benchislett for the original fix.

Purpose

FIX #32626
FIX #33450

Problem: TRTLLM attention requires that num_decode_tokens be divisible by num_requests. However, during DP we sometimes do a dummy run on one of the workers so they don't get out of sync: it such cases, we pad for attention but the attention metadata builder receives the padded number of requests and unpadded number of tokens.

This leads to a crash where we see num_decode_tokens==1 but num_requests==2. I tracked this down to _build_attention_metadata only seeing num_tokens=num_tokens_unpadded and having num_tokens_padded=None (omitted, default is None) even when pad_attn is True.

I'm not sure if this behaviour is intentional, or what the desired strategy for full-graph attention padding for DP will be. If num_reqs > num_decode_tokens is supposed to be supported, more investigation will be required to ensure that all backends (such as TRTLLM / FlashInfer) can handle this case. It might be as easy as broadening the check in the assert, or we may need to do some additional slicing and/or padding for some backends. I'm not too sure.

Testing

Ran this to launch the server on 4xB200:

vllm serve microsoft/Phi-mini-MoE-instruct -dp 2 -tp 2 --enable-expert-parallel --attention-backend flashinfer --attention-config.use_trtllm_attention true

and this to reproduce the crash:

vllm bench serve --random-input-len 128 --random-output-len 256 --random-range-ratio 0.5 --max-concurrency 64 --request-rate 5 --num-prompts 128 --num-warmups 0 --model $MODEL --base-url $BASE_URL

It doesn't trigger every single time, but according to my debugger it crashes consistently whenever the DP worker does a padded dummy run for a decode-only iteration (outside of startup).

With this patch, it finishes and the LM_Eval GSM8k score is unchanged (fewshot=5, non-chat-completions, ~80%). But I am not sure if this is safe more broadly or if it is omitted here on purpose.

[gemini-code-assist]

Code Review

This pull request addresses a crash during distributed dummy runs by correctly handling attention padding. The change in _dummy_run to pass num_tokens_padded to _build_attention_metadata when pad_attn is true is a direct and effective fix for the described issue. This change also aligns the logic with the execute_model function, improving code consistency. The fix is well-contained and correct. I approve of this change.