minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

fa3: build the topk>1 verify replay page table on-device (#31381)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-07-16 21:21:01 -07:00
committed by GitHub
parent 27ad9d11b1
commit 45824c69ca
1 changed files with 11 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/flashattention_backend.py
+11 -15
View File
@@ -2592,25 +2592,21 @@ class FlashAttentionBackend(AttentionBackend):
metadata = self.target_verify_metadata_topk_normal[bs]
metadata.cache_seqlens_int32.copy_(seq_lens)
# metadata.max_seq_len_q = self.speculative_num_draft_tokens, already set in capture
# Tight page-table bound when the CPU mirror is free.
metadata.max_seq_len_k = (
seq_lens_cpu.max().item()
if seq_lens_cpu is not None
else self.max_context_len
)
# Page table built on-device (self-guards on cache_seqlens).
# max_seq_len_k stays the static bound; its only replay reader
# is the SWA-merge capacity assert, capture-sized to match.
metadata.max_seq_len_k = self.max_context_len
# metadata.cu_seqlens_q already set in capture
metadata.cu_seqlens_k[1:].copy_(
torch.cumsum(metadata.cache_seqlens_int32, dim=0, dtype=torch.int32)
)
max_seq_pages = (
metadata.max_seq_len_k + self.page_size - 1
) // self.page_size
page_indices = self.req_to_token[
req_pool_indices[:, None],
self.decode_cuda_graph_metadata["strided_indices"][:max_seq_pages],
]
page_indices //= self.page_size
metadata.page_table[:, :max_seq_pages].copy_(page_indices)
build_trtllm_mha_page_table(
req_to_token=self.req_to_token,
req_pool_indices=req_pool_indices,
cache_seqlens=metadata.cache_seqlens_int32,
page_table=metadata.page_table,
page_size=self.page_size,
)
# 2. The second half of metadata for draft tokens (per_batch_num_tokens = topk)
metadata_expand = self.target_verify_metadata_topk_expand[bs]