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