[Fix] avoid stream sync and torch compile in prefill for fa3 backend (#4932)

2025-03-30 13:53:44 -07:00
parent 032f8faaab
commit e62d60fe6d
7 changed files with 30 additions and 35 deletions
@@ -797,7 +797,7 @@ class FlashInferMLAMultiStepDraftBackend:
                encoder_lens=None,
                forward_mode=ForwardMode.DECODE,
                spec_info=forward_batch.spec_info,
-                seq_lens_cpu=forward_batch.decode_seq_lens_cpu,
+                seq_lens_cpu=forward_batch.seq_lens_cpu,
            )

        self.common_template(forward_batch, self.cuda_graph_kv_indices, call_fn)