From e25810ae29058299b7bf845c7ed572f2474a1d85 Mon Sep 17 00:00:00 2001 From: Maximilien de Bayser Date: Sun, 1 Dec 2024 23:05:32 -0300 Subject: [PATCH] Fill TorchSDPAAttentionMetadata seq_lens_field for prefill (#10799) Signed-off-by: Max de Bayser --- vllm/attention/backends/torch_sdpa.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/vllm/attention/backends/torch_sdpa.py b/vllm/attention/backends/torch_sdpa.py index 16e044b618c40..dafa5bb56acda 100644 --- a/vllm/attention/backends/torch_sdpa.py +++ b/vllm/attention/backends/torch_sdpa.py @@ -341,7 +341,11 @@ def build(self, seq_lens: List[int], query_lens: List[int], ) else: block_tables = torch.tensor([]) - seq_lens_tensor = torch.tensor([]) + seq_lens_tensor = torch.tensor( + input_data.seq_lens[:input_data.num_prefills], + dtype=torch.int32, + device="cpu", + ) # For multi-modal models placeholder_index_maps = None