Fix Illegal Memory Access when fa3 + spec + topk + page_size > 1 (#15469)

Co-authored-by: Liangsheng Yin <lsyincs@gmail.com>

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

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
+                # NOTE: the max decode length is speculative_num_steps - 1 (one token always generated by draft extend)
-                expand_page_table = cache_loc[:, :decode_length].clone()
+                # and we leave one extra for last_page_len, which -> speculative_num_steps for the page table
-                strided_indices_expand = torch.arange(
+                expand_page_table = torch.zeros(
-                    0,
+                    forward_batch.batch_size * self.topk,
-                    decode_length,
+                    self.speculative_num_steps,
-                    self.page_size,
+                    dtype=torch.int32,
                     device=self.device,
                 )
-                last_page_lens_broadcast = expanded_last_page_lens.unsqueeze(-1).expand(
+                # shape: [bs, num_steps, topk] -> [bs x topk, num_steps]
-                    -1, expand_page_table.shape[1]
+                cache_loc = forward_batch.out_cache_loc.view(
+                    -1, self.speculative_num_steps
                 )
-                expand_page_table -= last_page_lens_broadcast
+                draft_decode_set_expand_metadata(
-                expand_page_table = (
+                    cache_seqlens_int32=self.forward_metadata_spec_decode_expand.cache_seqlens_int32,
-                    expand_page_table[:, strided_indices_expand] // self.page_size
+                    page_table=expand_page_table,
-                )
+                    last_page_lens=last_page_lens,
-                self.forward_metadata_spec_decode_expand.page_table = (
+                    decode_length=decode_length,
-                    expand_page_table.to(torch.int32)
+                    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.
+                        draft_decode_set_expand_metadata(
-                        strided_indices_expand = (
+                            cache_seqlens_int32=metadata_expand.cache_seqlens_int32,
-                            self.draft_decode_metadata_topk_expand.get(
+                            page_table=metadata_expand.page_table,
-                                "strided_indices_expand"
+                            last_page_lens=last_page_lens,
-                            )
+                            decode_length=decode_length,
-                        )
+                            cache_loc=cache_loc,
-                        update_draft_decode_set_expand_metadata_with_page_size(
+                            topk=self.topk,
-                            metadata_expand.cache_seqlens_int32,  # Modifies
+                            page_size=self.page_size,
-                            metadata_expand.page_table,  # Modifies
-                            cache_loc,
-                            last_page_lens,
-                            strided_indices_expand,
-                            decode_length,
-                            bs,
-                            self.topk,
-                            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()
+    cache_loc = (cache_loc // page_size).to(torch.int32)
-    last_page_lens_broadcast = expanded_last_page_lens.unsqueeze(-1).expand(
+    if cache_loc.dim() == 1:
-        -1, expand_page_table.shape[1]
+        cache_loc = cache_loc.unsqueeze(0)
-    )
+    # Vectorized torch.unique_consecutive: track value change points then scatter
-    num_seqs = expand_page_table.shape[0]
+    mask = torch.ones_like(cache_loc, dtype=torch.bool)
-    expand_page_table -= last_page_lens_broadcast[:num_seqs]
+    mask[:, 1:] = cache_loc[:, 1:] != cache_loc[:, :-1]
-    expand_page_table = (
+    positions = mask.cumsum(dim=1) - 1
-        expand_page_table[
+    num_seqs = cache_loc.shape[0]
-            :, strided_indices_expand[: (decode_length + page_size - 1) // page_size]
+    page_table[:num_seqs, :].scatter_(1, positions, cache_loc)
-        ]
-        // 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)

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,53 +369,100 @@ class TestUpdateDraftDecodeSetExpandMetadata(CustomTestCase):
                 ),
                 torch.tensor(
                     [
-                        [5],
+                        [5, 6],
-                        [7],
+                        [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],
+                        [27, 28, 29, 30, 31, 32],
-                        [35, 36],
+                        [35, 36, 37, 38, 39, 40],
                     ],
                     dtype=torch.int32,
                 ),
                 torch.tensor(
                     [
-                        [6],
+                        [6, 7, 8, 0, 0, 0],
-                        [8],
+                        [8, 9, 10, 0, 0, 0],
                     ],
                     dtype=torch.int32,
                 ),
+                5,
             ),
         ]
 
         last_page_lens = torch.tensor([3], dtype=torch.int32)
-        strided_indices_expand = torch.arange(
+        for cache_loc, expected_page_table, decode_length in cases:
-            0, decode_length, page_size, dtype=torch.long
-        )
-
-        for cache_loc, expected_page_table in cases:
             cache_seqlens_int32 = torch.zeros(bs * topk, dtype=torch.int32)
-            page_table = torch.zeros(bs * topk, decode_length, dtype=torch.int32)
+            page_table = torch.zeros_like(cache_loc, dtype=torch.int32)
-
+            draft_decode_set_expand_metadata(
-            update_draft_decode_set_expand_metadata_with_page_size(
                 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))
 

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,

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()