From 45824c69ca7e2248da17f660960ace65249907f5 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Thu, 16 Jul 2026 21:21:01 -0700 Subject: [PATCH] fa3: build the topk>1 verify replay page table on-device (#31381) --- .../attention/flashattention_backend.py | 26 ++++++++----------- 1 file changed, 11 insertions(+), 15 deletions(-) diff --git a/python/sglang/srt/layers/attention/flashattention_backend.py b/python/sglang/srt/layers/attention/flashattention_backend.py index 1230781be..a6c0122d5 100644 --- a/python/sglang/srt/layers/attention/flashattention_backend.py +++ b/python/sglang/srt/layers/attention/flashattention_backend.py @@ -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]