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
|
# creates a Dynamo shape guard. These graph modes never have empty
|
||||||
# batches.
|
# batches.
|
||||||
if not in_piecewise_or_breakable_cuda_graph:
|
if not in_piecewise_or_breakable_cuda_graph:
|
||||||
assert forward_batch.seq_lens_cpu is not None
|
if forward_batch.seq_lens.numel() == 0:
|
||||||
if len(forward_batch.seq_lens_cpu) == 0:
|
|
||||||
# this seems b/c max-pad, no worries?
|
# this seems b/c max-pad, no worries?
|
||||||
# if x.shape[0] != 0:
|
# if x.shape[0] != 0:
|
||||||
# print(
|
# print(
|
||||||
|
|||||||
@@ -755,6 +755,14 @@ class DeepseekSparseAttnBackend(
|
|||||||
page_table, repeats=self.speculative_num_draft_tokens, dim=0
|
page_table, repeats=self.speculative_num_draft_tokens, dim=0
|
||||||
)
|
)
|
||||||
elif forward_batch.forward_mode.is_draft_extend_v2():
|
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 (
|
assert (
|
||||||
forward_batch.extend_seq_lens_cpu is not None
|
forward_batch.extend_seq_lens_cpu is not None
|
||||||
and forward_batch.extend_seq_lens is not None
|
and forward_batch.extend_seq_lens is not None
|
||||||
|
|||||||
Reference in New Issue
Block a user