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

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

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