Unverified Commit e25810ae authored by Maximilien de Bayser's avatar Maximilien de Bayser Committed by GitHub
Browse files

Fill TorchSDPAAttentionMetadata seq_lens_field for prefill (#10799)


Signed-off-by: default avatarMax de Bayser <mbayser@br.ibm.com>
parent 073a4bd1
...@@ -341,7 +341,11 @@ class TorchSDPAMetadataBuilder(AttentionMetadataBuilder[TorchSDPAMetadata]): ...@@ -341,7 +341,11 @@ class TorchSDPAMetadataBuilder(AttentionMetadataBuilder[TorchSDPAMetadata]):
) )
else: else:
block_tables = torch.tensor([]) 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 # For multi-modal models
placeholder_index_maps = None placeholder_index_maps = None
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment