fix: avoid DSA indexer CPU seq lens fallback (#29798)

Co-authored-by: San Tian <san-tian@users.noreply.github.com>
This commit is contained in:
san-tian
2026-07-02 22:50:24 -07:00
committed by GitHub
co-authored by San Tian
parent e878c6ebdd
commit 70b6c06793
2 changed files with 9 additions and 2 deletions
@@ -1835,8 +1835,7 @@ class Indexer(MultiPlatformOp):
# creates a Dynamo shape guard. These graph modes never have empty
# batches.
if not in_piecewise_or_breakable_cuda_graph:
assert forward_batch.seq_lens_cpu is not None
if len(forward_batch.seq_lens_cpu) == 0:
if forward_batch.seq_lens.numel() == 0:
# this seems b/c max-pad, no worries?
# if x.shape[0] != 0:
# print(
@@ -755,6 +755,14 @@ class DeepseekSparseAttnBackend(
page_table, repeats=self.speculative_num_draft_tokens, dim=0
)
elif forward_batch.forward_mode.is_draft_extend_v2():
if forward_batch.extend_prefix_lens_cpu is None:
assert forward_batch.extend_prefix_lens is not None
forward_batch.extend_prefix_lens_cpu = (
forward_batch.extend_prefix_lens.cpu().tolist()
)
if forward_batch.seq_lens_cpu is None:
forward_batch.seq_lens_cpu = forward_batch.seq_lens.cpu()
forward_batch.seq_lens_sum = int(forward_batch.seq_lens_cpu.sum())
assert (
forward_batch.extend_seq_lens_cpu is not None
and forward_batch.extend_seq_lens is not None