diff --git a/python/sglang/srt/layers/attention/flashattention_backend.py b/python/sglang/srt/layers/attention/flashattention_backend.py index 98e0e9ecc..672605560 100644 --- a/python/sglang/srt/layers/attention/flashattention_backend.py +++ b/python/sglang/srt/layers/attention/flashattention_backend.py @@ -2519,7 +2519,8 @@ def update_draft_decode_set_expand_metadata_with_page_size( last_page_lens_broadcast = expanded_last_page_lens.unsqueeze(-1).expand( -1, expand_page_table.shape[1] ) - expand_page_table -= last_page_lens_broadcast + num_seqs = expand_page_table.shape[0] + expand_page_table -= last_page_lens_broadcast[:num_seqs] expand_page_table = ( expand_page_table[ :, strided_indices_expand[: (decode_length + page_size - 1) // page_size] @@ -2527,4 +2528,4 @@ def update_draft_decode_set_expand_metadata_with_page_size( // page_size ) max_seq_pages_expand = (decode_length + page_size - 1) // page_size - page_table[:, :max_seq_pages_expand].copy_(expand_page_table) + page_table[:num_seqs, :max_seq_pages_expand].copy_(expand_page_table) diff --git a/test/registered/spec/eagle/test_eagle_infer_a.py b/test/registered/spec/eagle/test_eagle_infer_a.py index 5c7e0d956..e70a92ece 100644 --- a/test/registered/spec/eagle/test_eagle_infer_a.py +++ b/test/registered/spec/eagle/test_eagle_infer_a.py @@ -228,6 +228,7 @@ class TestEAGLERadixCache(CustomTestCase): # Chunked prefill & Page Size > 1 {**self.BASE_CONFIG, "chunked_prefill_size": 64, "page_size": 4}, {**self.BASE_CONFIG, "page_size": 4}, + {**self.BASE_CONFIG, "cuda_graph_bs": [5], "page_size": 4}, # Preferred by some kernels {**self.BASE_CONFIG, "page_size": 64}, # Disable CUDA Graph