Support Page Size > 1 (when top k = 1) for FA3 Spec Decode

python/sglang/srt/layers/attention/flashattention_backend.py

@@ -767,9 +767,7 @@ def init_cuda_graph_state(self, max_bs: int):
             "cu_seqlens_q": torch.arange(
                 0, max_bs + 1, dtype=torch.int32, device=self.device
             ),
-            "cu_seqlens_k": torch.zeros(
-                max_bs + 1, dtype=torch.int32, device=self.device
-            ),
+            # cu_seqlens_k will be computed in capture_cuda_graph
             "page_table": torch.zeros(
                 max_bs,
                 (self.max_context_len + self.page_size - 1) // self.page_size,
@@ -789,9 +787,7 @@ def init_cuda_graph_state(self, max_bs: int):
 
         self.target_verify_metadata = {
             "cache_seqlens": torch.zeros(max_bs, dtype=torch.int32, device=self.device),
-            "cu_seqlens_q": torch.zeros(
-                max_bs + 1, dtype=torch.int32, device=self.device
-            ),
+            # cu_seqlens_q will be computed in capture_cuda_graph
             "cu_seqlens_k": torch.zeros(
                 max_bs + 1, dtype=torch.int32, device=self.device
             ),
@@ -958,11 +954,15 @@ def init_forward_metadata_replay_cuda_graph(
                     )
                 )
 
-                page_table = self.req_to_token[
-                    req_pool_indices, : metadata.max_seq_len_k
+                max_seq_pages = (
+                    metadata.max_seq_len_k + self.page_size - 1
+                ) // self.page_size
+                page_indices = self.req_to_token[
+                    :,
+                    self.decode_cuda_graph_metadata["strided_indices"][:max_seq_pages],
                 ]
-
-                metadata.page_table[:, : metadata.max_seq_len_k].copy_(page_table)
+                page_indices = page_indices[req_pool_indices] // self.page_size
+                metadata.page_table[:, :max_seq_pages].copy_(page_indices)
             else:
                 # Normal Decode
                 max_len = seq_lens_cpu.max().item()
@@ -984,7 +984,6 @@ def init_forward_metadata_replay_cuda_graph(
                 ]
                 page_indices //= self.page_size
                 metadata.page_table[:, :max_seq_pages].copy_(page_indices)
-                metadata.page_table[:, max_seq_pages:].fill_(0)
 
         elif forward_mode.is_target_verify():
             metadata = self.target_verify_metadata[bs]
@@ -1003,8 +1002,16 @@ def init_forward_metadata_replay_cuda_graph(
                     (1, 0),
                 )
             )
-            page_table = self.req_to_token[req_pool_indices, : metadata.max_seq_len_k]
-            metadata.page_table[:, : metadata.max_seq_len_k].copy_(page_table)
+
+            max_seq_pages = (
+                metadata.max_seq_len_k + self.page_size - 1
+            ) // self.page_size
+            page_indices = self.req_to_token[
+                :,
+                self.decode_cuda_graph_metadata["strided_indices"][:max_seq_pages],
+            ]
+            page_indices = page_indices[req_pool_indices] // self.page_size
+            metadata.page_table[:, :max_seq_pages].copy_(page_indices)
 
         if encoder_lens is not None:
             # Only support encoder size 1 for now