Fix tensor mismatch error in sepc + topk > 1 + page_size > 1 (#14874)

This commit is contained in:
Ziming Huang
2025-12-15 01:10:58 +08:00
committed by GitHub
parent 47633c191f
commit 997ea57eaf
2 changed files with 4 additions and 2 deletions

View File

@@ -2519,7 +2519,8 @@ def update_draft_decode_set_expand_metadata_with_page_size(
last_page_lens_broadcast = expanded_last_page_lens.unsqueeze(-1).expand(
-1, expand_page_table.shape[1]
)
expand_page_table -= last_page_lens_broadcast
num_seqs = expand_page_table.shape[0]
expand_page_table -= last_page_lens_broadcast[:num_seqs]
expand_page_table = (
expand_page_table[
:, strided_indices_expand[: (decode_length + page_size - 1) // page_size]
@@ -2527,4 +2528,4 @@ def update_draft_decode_set_expand_metadata_with_page_size(
// page_size
)
max_seq_pages_expand = (decode_length + page_size - 1) // page_size
page_table[:, :max_seq_pages_expand].copy_(expand_page_table)
page_table[:num_seqs, :max_seq_pages_expand].copy_(expand_page_table)