diff --git a/python/sglang/srt/layers/attention/flashattention_backend.py b/python/sglang/srt/layers/attention/flashattention_backend.py index 6c5d67220..0d41120d6 100644 --- a/python/sglang/srt/layers/attention/flashattention_backend.py +++ b/python/sglang/srt/layers/attention/flashattention_backend.py @@ -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: diff --git a/test/registered/8-gpu-models/test_mimo_models.py b/test/registered/8-gpu-models/test_mimo_models.py index 6baa23bef..dd111133b 100644 --- a/test/registered/8-gpu-models/test_mimo_models.py +++ b/test/registered/8-gpu-models/test_mimo_models.py @@ -27,6 +27,8 @@ class TestMiMoV2Flash(GSM8KMixin, SpecDecodingMixin, DefaultServerBase): "128", "--cuda-graph-max-bs", "64", + "--page-size", + "64", "--mem-fraction-static", "0.75", "--speculative-algorithm",