chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Fix FA3 swa spec pg_size > 1 (#20369)

Browse Source
This commit is contained in:
Ke Bao
2026-03-12 11:42:01 +08:00
committed by GitHub
parent 649d6f2bc8
commit ae7c2397b9
2 changed files with 62 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

60
python/sglang/srt/layers/attention/flashattention_backend.py
View File

@@ -1534,6 +1534,20 @@ class FlashAttentionBackend(AttentionBackend):
),
}
if self.use_sliding_window_kv_pool:
self.target_verify_metadata["swa_page_table"] = torch.zeros(
max_bs,
max_num_pages,
dtype=torch.int32,
device=self.device,
)
self.draft_extend_metadata["swa_page_table"] = torch.zeros(
max_bs,
max_num_pages,
dtype=torch.int32,
device=self.device,
)
if self.topk > 1:
self.target_verify_metadata_topk_normal = {
"cache_seqlens": torch.zeros(
@@ -1668,6 +1682,10 @@ class FlashAttentionBackend(AttentionBackend):
metadata.page_table = self.decode_cuda_graph_metadata[
"page_table_draft_decode"
][:bs, :]
if self.use_sliding_window_kv_pool:
metadata.swa_page_table = self.decode_cuda_graph_metadata[
"swa_page_table"
][:bs, :]
self.decode_cuda_graph_metadata[bs] = metadata
else:
# When top k > 1, we need two specific draft decode metadata, and then merge states
@@ -1764,6 +1782,11 @@ class FlashAttentionBackend(AttentionBackend):
metadata.page_table = self.target_verify_metadata["page_table"][:bs, :]
if self.use_sliding_window_kv_pool:
metadata.swa_page_table = self.target_verify_metadata[
"swa_page_table"
][:bs, :]
self.target_verify_metadata[bs] = metadata
else:
# When topk > 1, we need two specific target verify metadata, and then merge states
@@ -1848,6 +1871,11 @@ class FlashAttentionBackend(AttentionBackend):
]
metadata.page_table = self.draft_extend_metadata["page_table"][:bs, :]
if self.use_sliding_window_kv_pool:
metadata.swa_page_table = self.draft_extend_metadata["swa_page_table"][
:bs, :
]
self.draft_extend_metadata[bs] = metadata
if encoder_lens is not None:
@@ -1910,6 +1938,12 @@ class FlashAttentionBackend(AttentionBackend):
seq_lens,
self.speculative_step_id + 1,
self.page_size,
metadata.swa_page_table,
(
self.token_to_kv_pool
if self.use_sliding_window_kv_pool
else None
),
)
else:
@@ -2006,6 +2040,18 @@ class FlashAttentionBackend(AttentionBackend):
req_pool_indices[:, None],
self.decode_cuda_graph_metadata["strided_indices"][:max_seq_pages],
]
if (
self.use_sliding_window_kv_pool
and metadata.swa_page_table is not None
):
swa_page_indices = (
self.token_to_kv_pool.translate_loc_from_full_to_swa(
page_indices
)
)
metadata.swa_page_table[:, :max_seq_pages].copy_(
swa_page_indices // self.page_size
)
page_indices //= self.page_size
metadata.page_table[:, :max_seq_pages].copy_(page_indices)
else:
@@ -2121,6 +2167,13 @@ class FlashAttentionBackend(AttentionBackend):
req_pool_indices[:, None],
self.draft_extend_metadata["strided_indices"][:max_seq_pages],
]
if self.use_sliding_window_kv_pool and metadata.swa_page_table is not None:
swa_page_indices = self.token_to_kv_pool.translate_loc_from_full_to_swa(
page_indices
)
metadata.swa_page_table[:, :max_seq_pages].copy_(
swa_page_indices // self.page_size
)
metadata.page_table[:, :max_seq_pages].copy_(page_indices // self.page_size)
elif forward_mode.is_draft_extend_v2():
@@ -2169,6 +2222,13 @@ class FlashAttentionBackend(AttentionBackend):
req_pool_indices[:, None],
self.draft_extend_metadata["strided_indices"][:max_seq_pages],
]
if self.use_sliding_window_kv_pool and metadata.swa_page_table is not None:
swa_page_indices = self.token_to_kv_pool.translate_loc_from_full_to_swa(
page_indices
)
metadata.swa_page_table[:, :max_seq_pages].copy_(
swa_page_indices // self.page_size
)
metadata.page_table[:, :max_seq_pages].copy_(page_indices // self.page_size)
if encoder_lens is not None: