[Fix] Unify ForwardBatch extend lens cpu fields to their declared list type (#30896)
This commit is contained in:
@@ -222,15 +222,11 @@ class TestFlashAttentionBackend(CustomTestCase):
|
||||
extend_prefix_lens=torch.tensor(
|
||||
[prefix_len] * self.batch_size, device=self.device
|
||||
),
|
||||
extend_prefix_lens_cpu=torch.tensor(
|
||||
[prefix_len] * self.batch_size, device="cpu"
|
||||
),
|
||||
extend_prefix_lens_cpu=[prefix_len] * self.batch_size,
|
||||
extend_seq_lens=torch.tensor(
|
||||
[q_len] * self.batch_size, device=self.device
|
||||
),
|
||||
extend_seq_lens_cpu=torch.tensor(
|
||||
[q_len] * self.batch_size, device="cpu"
|
||||
),
|
||||
extend_seq_lens_cpu=[q_len] * self.batch_size,
|
||||
)
|
||||
if attn_cp_size > 1:
|
||||
forward_batch.attn_cp_metadata = type(
|
||||
|
||||
@@ -190,15 +190,11 @@ class TestFlashAttentionMLABackend(CustomTestCase):
|
||||
extend_prefix_lens=torch.tensor(
|
||||
[prefix_len] * self.batch_size, device=self.device
|
||||
),
|
||||
extend_prefix_lens_cpu=torch.tensor(
|
||||
[prefix_len] * self.batch_size, device="cpu"
|
||||
),
|
||||
extend_prefix_lens_cpu=[prefix_len] * self.batch_size,
|
||||
extend_seq_lens=torch.tensor(
|
||||
[q_len] * self.batch_size, device=self.device
|
||||
),
|
||||
extend_seq_lens_cpu=torch.tensor(
|
||||
[q_len] * self.batch_size, device="cpu"
|
||||
),
|
||||
extend_seq_lens_cpu=[q_len] * self.batch_size,
|
||||
)
|
||||
|
||||
else: # ForwardMode.DECODE
|
||||
|
||||
Reference in New Issue
Block a user