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 = self.target_verify_metadata_topk_normal[bs]
metadata.cache_seqlens_int32.copy_(seq_lens) metadata.cache_seqlens_int32.copy_(seq_lens)
# metadata.max_seq_len_q = self.speculative_num_draft_tokens, already set in capture # metadata.max_seq_len_q = self.speculative_num_draft_tokens, already set in capture
# Tight page-table bound when the CPU mirror is free. # Page table built on-device (self-guards on cache_seqlens).
metadata.max_seq_len_k = ( # max_seq_len_k stays the static bound; its only replay reader
seq_lens_cpu.max().item() # is the SWA-merge capacity assert, capture-sized to match.
if seq_lens_cpu is not None metadata.max_seq_len_k = self.max_context_len
else self.max_context_len
)
# metadata.cu_seqlens_q already set in capture # metadata.cu_seqlens_q already set in capture
metadata.cu_seqlens_k[1:].copy_( metadata.cu_seqlens_k[1:].copy_(
torch.cumsum(metadata.cache_seqlens_int32, dim=0, dtype=torch.int32) torch.cumsum(metadata.cache_seqlens_int32, dim=0, dtype=torch.int32)
) )
max_seq_pages = ( build_trtllm_mha_page_table(
metadata.max_seq_len_k + self.page_size - 1 req_to_token=self.req_to_token,
) // self.page_size req_pool_indices=req_pool_indices,
page_indices = self.req_to_token[ cache_seqlens=metadata.cache_seqlens_int32,
req_pool_indices[:, None], page_table=metadata.page_table,
self.decode_cuda_graph_metadata["strided_indices"][:max_seq_pages], page_size=self.page_size,
] )
page_indices //= self.page_size
metadata.page_table[:, :max_seq_pages].copy_(page_indices)
# 2. The second half of metadata for draft tokens (per_batch_num_tokens = topk) # 2. The second half of metadata for draft tokens (per_batch_num_tokens = topk)
metadata_expand = self.target_verify_metadata_topk_expand[bs] metadata_expand = self.target_verify_metadata_topk_expand[bs]