[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
|
metadata.local_attn_metadata = None
|
||||||
return
|
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()
|
cu_seqlens_q_np = cu_seqlens_q.cpu().numpy()
|
||||||
seq_lens_np = cache_seqlens_int32.cpu().numpy()
|
seq_lens_np = cache_seqlens_int32.cpu().numpy()
|
||||||
(
|
(
|
||||||
|
|||||||
Reference in New Issue
Block a user