From 762846531f9805049aecc98027a6dc49ce33ba5c Mon Sep 17 00:00:00 2001 From: Yubo Wang Date: Tue, 23 Dec 2025 08:13:57 -0800 Subject: [PATCH] Fix Illegal Memory Access when fa3 + spec + topk + page_size > 1 (#15469) Co-authored-by: Liangsheng Yin --- benchmark/mtbench/bench_sglang_eagle.py | 4 +- .../attention/flashattention_backend.py | 103 +++++++----------- .../test/attention/test_flashattn_backend.py | 96 ++++++++++++---- .../spec/eagle/test_eagle_infer_a.py | 4 +- .../spec/eagle/test_eagle_infer_b.py | 15 +++ 5 files changed, 133 insertions(+), 89 deletions(-) diff --git a/benchmark/mtbench/bench_sglang_eagle.py b/benchmark/mtbench/bench_sglang_eagle.py index 7ad3dcb12..3eb6036c7 100644 --- a/benchmark/mtbench/bench_sglang_eagle.py +++ b/benchmark/mtbench/bench_sglang_eagle.py @@ -38,7 +38,7 @@ def write_answers(filename, model_id, questions, answers): "model_id": model_id, "choices": { "index": 0, - "prompt": [answers[i][0], answers[i][1]], + "turns": [answers[i][0], answers[i][1]], }, "tstamp": time.time(), } @@ -60,7 +60,7 @@ def main(args): # Construct prompts questions = load_questions(args.question_file)[: args.num_questions] arguments = [ - {"question_1": q["prompt"][0], "question_2": q["prompt"][1]} for q in questions + {"question_1": q["turns"][0], "question_2": q["turns"][1]} for q in questions ] # Select backend diff --git a/python/sglang/srt/layers/attention/flashattention_backend.py b/python/sglang/srt/layers/attention/flashattention_backend.py index ed4ec1273..7402ce516 100644 --- a/python/sglang/srt/layers/attention/flashattention_backend.py +++ b/python/sglang/srt/layers/attention/flashattention_backend.py @@ -668,24 +668,28 @@ class FlashAttentionBackend(AttentionBackend): self.forward_metadata_spec_decode_expand.cache_seqlens_int32 += ( expanded_last_page_lens ) - decode_length = self.speculative_step_id + 1 - expand_page_table = cache_loc[:, :decode_length].clone() - strided_indices_expand = torch.arange( - 0, - decode_length, - self.page_size, + # NOTE: the max decode length is speculative_num_steps - 1 (one token always generated by draft extend) + # and we leave one extra for last_page_len, which -> speculative_num_steps for the page table + expand_page_table = torch.zeros( + forward_batch.batch_size * self.topk, + self.speculative_num_steps, + dtype=torch.int32, device=self.device, ) - last_page_lens_broadcast = expanded_last_page_lens.unsqueeze(-1).expand( - -1, expand_page_table.shape[1] + # shape: [bs, num_steps, topk] -> [bs x topk, num_steps] + cache_loc = forward_batch.out_cache_loc.view( + -1, self.speculative_num_steps ) - expand_page_table -= last_page_lens_broadcast - expand_page_table = ( - expand_page_table[:, strided_indices_expand] // self.page_size - ) - self.forward_metadata_spec_decode_expand.page_table = ( - expand_page_table.to(torch.int32) + draft_decode_set_expand_metadata( + cache_seqlens_int32=self.forward_metadata_spec_decode_expand.cache_seqlens_int32, + page_table=expand_page_table, + last_page_lens=last_page_lens, + decode_length=decode_length, + cache_loc=cache_loc, + topk=self.topk, + page_size=self.page_size, ) + self.forward_metadata_spec_decode_expand.page_table = expand_page_table self.forward_metadata = metadata @@ -1404,23 +1408,12 @@ class FlashAttentionBackend(AttentionBackend): ), "page_table": torch.zeros( max_bs * self.topk, - decode_length, + decode_length + 1, # Additional page for last partial page dtype=torch.int32, device=self.device, ), } - if self.page_size > 1: - # Used for indicing expand page_table - self.draft_decode_metadata_topk_expand["strided_indices_expand"] = ( - torch.arange( - 0, - self.speculative_num_steps, - self.page_size, - device=self.device, - ) - ) - if ( self.speculative_num_draft_tokens is not None and self.speculative_num_draft_tokens > 0 @@ -1859,8 +1852,6 @@ class FlashAttentionBackend(AttentionBackend): # First attention handles seq_lens - last_page_lens if page size > 1. last_page_lens = seq_lens % self.page_size seq_lens = seq_lens - last_page_lens - # last_page_lens_cpu = last_page_lens.max().item() - # seq_lens_cpu -= last_page_lens_cpu metadata.cache_seqlens_int32.copy_(seq_lens) # metadata.max_seq_len_q = self.topk, already set in capture # metadata.cu_seqlens_q already set in capture @@ -1886,25 +1877,18 @@ class FlashAttentionBackend(AttentionBackend): # shape: [bs, num_steps, topk] -> [bs x topk, num_steps] cache_loc = out_cache_loc.view(-1, self.speculative_num_steps) if self.page_size > 1: - # Second attention handles last_page_len + decode part. - strided_indices_expand = ( - self.draft_decode_metadata_topk_expand.get( - "strided_indices_expand" - ) - ) - update_draft_decode_set_expand_metadata_with_page_size( - metadata_expand.cache_seqlens_int32, # Modifies - metadata_expand.page_table, # Modifies - cache_loc, - last_page_lens, - strided_indices_expand, - decode_length, - bs, - self.topk, - self.page_size, + draft_decode_set_expand_metadata( + cache_seqlens_int32=metadata_expand.cache_seqlens_int32, + page_table=metadata_expand.page_table, + last_page_lens=last_page_lens, + decode_length=decode_length, + cache_loc=cache_loc, + topk=self.topk, + page_size=self.page_size, ) else: - metadata_expand.page_table[: cache_loc.shape[0]].copy_( + num_seqs = cache_loc.shape[0] + metadata_expand.page_table[:num_seqs, :decode_length].copy_( cache_loc[:, :decode_length] ) # TODO: Handle local attention metadata for draft decode when llama4 eagle is supported @@ -2571,30 +2555,23 @@ def normal_decode_set_metadata( @torch.compile(dynamic=True, backend=get_compiler_backend()) -def update_draft_decode_set_expand_metadata_with_page_size( +def draft_decode_set_expand_metadata( cache_seqlens_int32: torch.Tensor, # Modifies page_table: torch.Tensor, # Modifies - cache_loc: torch.Tensor, last_page_lens: torch.Tensor, - strided_indices_expand: torch.Tensor, decode_length: int, - bs: int, + cache_loc: torch.Tensor, topk: int, page_size: int, ): expanded_last_page_lens = last_page_lens.repeat_interleave(topk) cache_seqlens_int32.copy_(decode_length + expanded_last_page_lens) - expand_page_table = cache_loc[:, :decode_length].clone() - last_page_lens_broadcast = expanded_last_page_lens.unsqueeze(-1).expand( - -1, expand_page_table.shape[1] - ) - 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] - ] - // page_size - ) - max_seq_pages_expand = (decode_length + page_size - 1) // page_size - page_table[:num_seqs, :max_seq_pages_expand].copy_(expand_page_table) + cache_loc = (cache_loc // page_size).to(torch.int32) + if cache_loc.dim() == 1: + cache_loc = cache_loc.unsqueeze(0) + # Vectorized torch.unique_consecutive: track value change points then scatter + mask = torch.ones_like(cache_loc, dtype=torch.bool) + mask[:, 1:] = cache_loc[:, 1:] != cache_loc[:, :-1] + positions = mask.cumsum(dim=1) - 1 + num_seqs = cache_loc.shape[0] + page_table[:num_seqs, :].scatter_(1, positions, cache_loc) diff --git a/python/sglang/test/attention/test_flashattn_backend.py b/python/sglang/test/attention/test_flashattn_backend.py index a58e19498..a134d18a7 100644 --- a/python/sglang/test/attention/test_flashattn_backend.py +++ b/python/sglang/test/attention/test_flashattn_backend.py @@ -5,7 +5,7 @@ import torch from sglang.srt.configs.model_config import AttentionArch from sglang.srt.layers.attention.flashattention_backend import ( FlashAttentionBackend, - update_draft_decode_set_expand_metadata_with_page_size, + draft_decode_set_expand_metadata, ) from sglang.srt.layers.attention.torch_native_backend import TorchNativeAttnBackend from sglang.srt.layers.radix_attention import RadixAttention @@ -350,8 +350,13 @@ class TestFlashAttentionBackend(CustomTestCase): class TestUpdateDraftDecodeSetExpandMetadata(CustomTestCase): - def test_update_draft_decode_set_expand_metadata_with_page_size(self): - bs, topk, decode_length, page_size = 1, 2, 1, 4 + """ + All the test cases examples have 1 additional cache location than the decode length. + This is to align with the current allocation logic. It does not affect the correctness. + """ + + def test_draft_decode_set_expand_metadata(self): + bs, topk, page_size = 1, 2, 4 cases = [ ( @@ -364,56 +369,103 @@ class TestUpdateDraftDecodeSetExpandMetadata(CustomTestCase): ), torch.tensor( [ - [5], - [7], + [5, 6], + [7, 8], ], dtype=torch.int32, ), + 1, ), + # Decode span multiple pages: + # duplicated kv cache: 24, 25, 26 + # decode locations: 27, 28, 29, 30, 31, 32 + # We need 3 pages in total. ( torch.tensor( [ - [27, 28], - [35, 36], + [27, 28, 29, 30, 31, 32], + [35, 36, 37, 38, 39, 40], ], dtype=torch.int32, ), torch.tensor( [ - [6], - [8], + [6, 7, 8, 0, 0, 0], + [8, 9, 10, 0, 0, 0], ], dtype=torch.int32, ), + 5, ), ] last_page_lens = torch.tensor([3], dtype=torch.int32) - strided_indices_expand = torch.arange( - 0, decode_length, page_size, dtype=torch.long - ) - - for cache_loc, expected_page_table in cases: + for cache_loc, expected_page_table, decode_length in cases: cache_seqlens_int32 = torch.zeros(bs * topk, dtype=torch.int32) - page_table = torch.zeros(bs * topk, decode_length, dtype=torch.int32) - - update_draft_decode_set_expand_metadata_with_page_size( + page_table = torch.zeros_like(cache_loc, dtype=torch.int32) + draft_decode_set_expand_metadata( cache_seqlens_int32=cache_seqlens_int32, page_table=page_table, - cache_loc=cache_loc, last_page_lens=last_page_lens, - strided_indices_expand=strided_indices_expand, decode_length=decode_length, - bs=bs, + cache_loc=cache_loc, topk=topk, page_size=page_size, ) - expected_cache_seqlens = torch.tensor([4, 4], dtype=torch.int32) - + expected_cache_seqlens = torch.tensor( + [decode_length + 3, decode_length + 3], dtype=torch.int32 + ) self.assertTrue(torch.equal(cache_seqlens_int32, expected_cache_seqlens)) self.assertTrue(torch.equal(page_table, expected_page_table)) + def test_update_draft_decode_set_expand_metadata_multi_batch(self): + """ + Ensure expand metadata works when batch size > 1 and last pages differ. + """ + bs, topk, decode_length, page_size = 3, 2, 3, 4 + cache_loc = torch.tensor( + [ + # First batch: last page duplicate is 1, consecutive pages + [1, 2, 3, 4], + [6, 7, 8, 9], + # Second batch: last page duplicate is 3, non-consecutive pages + [3, 8, 9, 10], + [14, 15, 16, 17], + # Third batch: last page duplicate is 0, consecutive pages + [0, 1, 2, 3], + [4, 5, 6, 7], + ], + dtype=torch.int32, + ) + cache_seqlens_int32 = torch.zeros(bs * topk, dtype=torch.int32) + last_page_lens = torch.tensor([1, 3, 0], dtype=torch.int32) + page_table = torch.zeros_like(cache_loc, dtype=torch.int32) + draft_decode_set_expand_metadata( + cache_seqlens_int32=cache_seqlens_int32, + page_table=page_table, + last_page_lens=last_page_lens, + decode_length=decode_length, + cache_loc=cache_loc, + topk=topk, + page_size=page_size, + ) + + expected_cache_seqlens = torch.tensor([4, 4, 6, 6, 3, 3], dtype=torch.int32) + expected_page_table = torch.tensor( + [ + [0, 1, 0, 0], + [1, 2, 0, 0], + [0, 2, 0, 0], + [3, 4, 0, 0], + [0, 0, 0, 0], + [1, 0, 0, 0], + ], + dtype=torch.int32, + ) + self.assertTrue(torch.equal(cache_seqlens_int32, expected_cache_seqlens)) + self.assertTrue(torch.equal(page_table, expected_page_table)) + if __name__ == "__main__": unittest.main() diff --git a/test/registered/spec/eagle/test_eagle_infer_a.py b/test/registered/spec/eagle/test_eagle_infer_a.py index 270dbbbd9..9a9bdd3af 100644 --- a/test/registered/spec/eagle/test_eagle_infer_a.py +++ b/test/registered/spec/eagle/test_eagle_infer_a.py @@ -228,9 +228,9 @@ class TestEAGLERadixCache(CustomTestCase): # Chunked prefill & Page Size > 1 {**self.BASE_CONFIG, "chunked_prefill_size": 64, "page_size": 4}, {**self.BASE_CONFIG, "page_size": 4}, + # Large page size tend to expose IMA bugs. + {**self.BASE_CONFIG, "page_size": 256}, {**self.BASE_CONFIG, "cuda_graph_bs": [5], "page_size": 4}, - # Preferred by some kernels - {**self.BASE_CONFIG, "page_size": 64}, # Disable CUDA Graph { **self.BASE_CONFIG, diff --git a/test/registered/spec/eagle/test_eagle_infer_b.py b/test/registered/spec/eagle/test_eagle_infer_b.py index e4d3ebf40..ca49db900 100644 --- a/test/registered/spec/eagle/test_eagle_infer_b.py +++ b/test/registered/spec/eagle/test_eagle_infer_b.py @@ -333,5 +333,20 @@ class TestEAGLEServerPageSizeTopk(TestEAGLEServerBasic): ] +class TestEAGLEServerPageSizeTopkFA3(TestEAGLEServerBasic): + # default topk=8 and tokens=64 + spec_topk = 5 + spec_steps = 8 + spec_tokens = 64 + + extra_args = [ + "--page-size=256", + "--attention-backend=fa3", + "--cuda-graph-max-bs=5", + "--dtype=float16", + "--max-running-requests=8", + ] + + if __name__ == "__main__": unittest.main()