Fill TorchSDPAAttentionMetadata seq_lens_field for prefill (#10799)

Signed-off-by: Max de Bayser <mbayser@br.ibm.com>

Fill TorchSDPAAttentionMetadata seq_lens_field for prefill (#10799)
Signed-off-by: Max de Bayser <mbayser@br.ibm.com>
e25810ae · Maximilien de Bayser · GitHub · 073a4bd1 · e25810ae
Unverified Commit e25810ae authored Dec 01, 2024 by Maximilien de Bayser Committed by GitHub Dec 02, 2024
Hide whitespace changes
Inline Side-by-side

Showing with 5 additions and 1 deletion

vllm/attention/backends/torch_sdpa.py vllm/attention/backends/torch_sdpa.py +5 -1

No files found.
--- a/vllm/attention/backends/torch_sdpa.py
+++ b/vllm/attention/backends/torch_sdpa.py
@@ -341,7 +341,11 @@ class TorchSDPAMetadataBuilder(AttentionMetadataBuilder[TorchSDPAMetadata]):
            )
        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