[Bugfix] Exclude DP token padding from draft attention metadata (#55458)
Signed-off-by: khluu <khluu000@gmail.com> Signed-off-by: Kevin H. Luu <khluu000@gmail.com> Co-authored-by: Codex <noreply@openai.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Kevin H. Luu committed
Sep 9, 2026 at 01:47 UTC
5acd95906b7c7a54dde89396d2bb06fe28ebeed0
1 file changed
+4
-2
vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py
+4
-2
@@ -511,7 +511,8 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
511
attn_metadata = self._build_draft_attn_metadata(
512
num_reqs=num_reqs,
513
num_reqs_padded=batch_desc.num_reqs or num_reqs,
514
- num_tokens_padded=batch_desc.num_tokens,
514
+ # One query per request; exclude DP-only model padding.
515
+ num_tokens_padded=batch_desc.num_reqs or num_reqs,
516
seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,
517
step=step,
518
)
@@ -559,7 +560,8 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
560
attn_metadata = self._build_draft_attn_metadata(
561
num_reqs=num_reqs,
562
num_reqs_padded=batch_desc.num_reqs or num_reqs,
562
- num_tokens_padded=batch_desc.num_tokens,
563
+ # One query per request; exclude DP-only model padding.
564
+ num_tokens_padded=batch_desc.num_reqs or num_reqs,
565
seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,
566
step=1,
567
)