Merge 8756b5ed1568cf3676d7be5d39d022c11a36fbad into 254f6b986720c92ddf97fbb1a6a6465da8e87e29

This commit is contained in:
Kevin McKay 2025-12-25 00:07:00 +00:00 committed by GitHub
commit 143590cdc6
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -122,7 +122,11 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
).unsqueeze(0) < seq_lens_device.unsqueeze(1)
paged_kv_indices = block_table_tensor[mask]
paged_kv_last_page_len = torch.where(seq_lens_device == 0, 1, seq_lens_device)
# kernel block size is always 1, so each page has exactly 1 token.
# last_page_len should always be 1 regardless of sequence length.
paged_kv_last_page_len = torch.ones(
num_reqs, dtype=seq_lens_device.dtype, device=device
)
paged_kv_indptr = torch.cat(
[