From a09e85d677f6a28d1552b303c91080edfc670525 Mon Sep 17 00:00:00 2001
From: Zhiy-Zhang <44971983+Zhiy-Zhang@users.noreply.github.com>
Date: Tue, 9 Jun 2026 08:14:29 +0800
Subject: [PATCH] Fix(spec): Fix the crash issue in the FA3 backend when
 running with top-k > 1 and page_size > 1 (#25077)

Co-authored-by: qiufan.zzy <qiufan.zzy@antgroup.com>
---
 .../attention/test_flashattn_backend.py       | 31 +++++++++++++++++++
 1 file changed, 31 insertions(+)

diff --git a/test/manual/attention/test_flashattn_backend.py b/test/manual/attention/test_flashattn_backend.py
index 466871df6..5b99abceb 100644
--- a/test/manual/attention/test_flashattn_backend.py
+++ b/test/manual/attention/test_flashattn_backend.py
@@ -509,6 +509,37 @@ class TestUpdateDraftDecodeSetExpandMetadata(CustomTestCase):
             self.assertTrue(torch.equal(cache_seqlens_int32, expected_cache_seqlens))
             self.assertTrue(torch.equal(page_table, expected_page_table))
 
+    def test_draft_decode_set_expand_metadata_page_crossing(self):
+        """
+        Regression for fa3 EAGLE draft decode with topk > 1 and page_size > 1.
+        cache_loc arrives num_steps-wide; callers pre-slice it to `decode_length`
+        (the live draft tokens) before this helper runs, so the dedup'd scatter
+        never writes past the (decode_length + 1)-wide expand page_table row even
+        when consecutive draft tokens land on distinct pages.
+        """
+        bs, topk, page_size = 1, 2, 4
+        decode_length = 2
+        last_page_lens = torch.tensor([3], dtype=torch.int32)
+        # 2 live draft tokens per (batch, topk) crossing into distinct pages.
+        cache_loc = torch.tensor([[23, 28], [31, 36]], dtype=torch.int32)
+        cache_seqlens_int32 = torch.zeros(bs * topk, dtype=torch.int32)
+        # page_table is (decode_length + 1) wide (extra slot for the last partial
+        # page); the trailing column must stay zero.
+        page_table = torch.zeros(bs * topk, decode_length + 1, 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_page_table = torch.tensor([[5, 7, 0], [7, 9, 0]], dtype=torch.int32)
+        expected_cache_seqlens = torch.tensor([5, 5], dtype=torch.int32)
+        self.assertTrue(torch.equal(page_table, expected_page_table))
+        self.assertTrue(torch.equal(cache_seqlens_int32, expected_cache_seqlens))
+
     def test_update_draft_decode_set_expand_metadata_multi_batch(self):
         """
         Ensure expand metadata works when batch size > 1 and last pages differ.