[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>
This commit is contained in:
Kevin H. Luu
2026-09-09 01:47:45 -07:00
committed by GitHub
co-authored by Codex Wentao Ye mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
parent 114abd1c11
commit 5acd95906b
@@ -511,7 +511,8 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
attn_metadata = self._build_draft_attn_metadata(
num_reqs=num_reqs,
num_reqs_padded=batch_desc.num_reqs or num_reqs,
num_tokens_padded=batch_desc.num_tokens,
# One query per request; exclude DP-only model padding.
num_tokens_padded=batch_desc.num_reqs or num_reqs,
seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,
step=step,
)
@@ -559,7 +560,8 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
attn_metadata = self._build_draft_attn_metadata(
num_reqs=num_reqs,
num_reqs_padded=batch_desc.num_reqs or num_reqs,
num_tokens_padded=batch_desc.num_tokens,
# One query per request; exclude DP-only model padding.
num_tokens_padded=batch_desc.num_reqs or num_reqs,
seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,
step=1,
)