[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:
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,
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user