[Intel GPU] Fix incorrect KV-cache page table for local attention when page_size > 1 (#23757)
This commit is contained in:
@@ -925,6 +925,17 @@ class XPUAttentionBackend(AttentionBackend):
|
||||
metadata.local_attn_metadata = None
|
||||
return
|
||||
|
||||
# make_local_attention_virtual_batches expects a page-granularity block table:
|
||||
# column p is the logical page number, and the value stored at that column is the
|
||||
# physical page index. The raw req_to_token table is token-granularity (column i =
|
||||
# the KV slot for token i), so when page_size > 1 we must stride and divide first
|
||||
# so that block_starts = k_seqstarts_absolute // page_size correctly indexes the table.
|
||||
if self.page_size > 1:
|
||||
strided_indices = torch.arange(
|
||||
0, page_table.shape[1], self.page_size, device=page_table.device
|
||||
)
|
||||
page_table = page_table[:, strided_indices] // self.page_size
|
||||
|
||||
cu_seqlens_q_np = cu_seqlens_q.cpu().numpy()
|
||||
seq_lens_np = cache_seqlens_int32.cpu().numpy()
|
||||
(
|
||||
|
||||
Reference in New Issue
Block a user