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

This commit is contained in:
Baizhou Zhang
2025-03-30 13:53:44 -07:00
committed by GitHub
parent 032f8faaab
commit e62d60fe6d
7 changed files with 30 additions and 35 deletions

View File

@@ -491,10 +491,10 @@ class CudaGraphRunner:
self.seq_lens[:raw_bs].copy_(forward_batch.seq_lens)
self.out_cache_loc[:raw_num_token].copy_(forward_batch.out_cache_loc)
self.positions[:raw_num_token].copy_(forward_batch.positions)
if forward_batch.decode_seq_lens_cpu is not None:
if forward_batch.seq_lens_cpu is not None:
if bs != raw_bs:
self.seq_lens_cpu.fill_(1)
self.seq_lens_cpu[:raw_bs].copy_(forward_batch.decode_seq_lens_cpu)
self.seq_lens_cpu[:raw_bs].copy_(forward_batch.seq_lens_cpu)
if self.is_encoder_decoder:
self.encoder_lens[:raw_bs].copy_(forward_batch.encoder_lens)