[Spec] Add sync-free fast_prefill_plan for EAGLE draft-extend CUDA graph (#28854)

2026-06-22 15:15:15 -07:00
parent c0198fc277
commit 770d6b2825
2 changed files with 351 additions and 0 deletions
@@ -0,0 +1,187 @@
+"""Equivalence tests for the sync-free `fast_prefill_plan`.
+
+`fast_prefill_plan` replaces FlashInfer's `BatchPrefillWithPagedKVCacheWrapper.plan`
+in the EAGLE draft-extend CUDA graph: upstream plan() does blocking `.to("cpu")`
+copies to build host scheduling metadata, while fast_prefill_plan takes that
+metadata as host-known args and reaches `_cached_module.plan` with no readback.
+
+Correctness is proven end-to-end: the same draft-extend attention, planned two
+independent ways (upstream plan() vs fast_prefill_plan), must yield the SAME
+`run()` output on identical q/kv. A mutation check reverses the kv_indices handed
+to fast_prefill_plan and asserts the output DIVERGES, so we know the output is
+sensitive to the metadata under test and the equivalence is not vacuous.
+"""
+
+import unittest
+
+import torch
+
+from sglang.srt.layers.attention.flashinfer_backend import fast_prefill_plan
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+try:
+    from flashinfer import BatchPrefillWithPagedKVCacheWrapper
+
+    _HAS_FLASHINFER = True
+except ImportError:
+    _HAS_FLASHINFER = False
+
+register_cuda_ci(est_time=30, stage="base-b", runner_config="1-gpu-small")
+
+# Draft-extend layout: constant qo (num_tokens_per_req per req), page_size 1.
+# Non-uniform seq_lens make cumsum non-trivial, so a wrong per-row kv split is
+# caught instead of hidden by equal lengths.
+NUM_TOKENS_PER_REQ = 8
+SEQ_LENS = [37, 12, 89]
+NUM_QO_HEADS = 8
+NUM_KV_HEADS = 8
+HEAD_DIM = 128
+DTYPE = torch.float16
+
+
+@unittest.skipUnless(_HAS_FLASHINFER, "requires flashinfer")
+class TestFastPrefillPlan(CustomTestCase):
+    def setUp(self):
+        torch.manual_seed(0)
+        self.device = "cuda"
+        bs = len(SEQ_LENS)
+        self.bs = bs
+        seq_lens = torch.tensor(SEQ_LENS, dtype=torch.int32, device=self.device)
+
+        # Device inputs in the exact layout draft-extend feeds plan().
+        self.qo_indptr = torch.arange(
+            0,
+            (bs + 1) * NUM_TOKENS_PER_REQ,
+            step=NUM_TOKENS_PER_REQ,
+            dtype=torch.int32,
+            device=self.device,
+        )
+        self.kv_indptr = torch.zeros(bs + 1, dtype=torch.int32, device=self.device)
+        self.kv_indptr[1:] = torch.cumsum(seq_lens, dim=0)
+        self.total_kv = int(self.kv_indptr[-1].item())
+        self.total_q = int(self.qo_indptr[-1].item())
+        self.kv_indices = torch.arange(
+            self.total_kv, dtype=torch.int32, device=self.device
+        )
+        self.last_page_len = torch.ones(bs, dtype=torch.int32, device=self.device)
+
+        # Host metadata the fast path is handed (page_size==1 -> token-level).
+        seq_lens_cpu = seq_lens.cpu()
+        self.qo_indptr_host = self.qo_indptr.cpu()
+        self.kv_indptr_host = self.kv_indptr.cpu()
+        self.kv_lens_host = seq_lens_cpu
+        self.max_q_len = NUM_TOKENS_PER_REQ
+        self.max_kv_len = int(seq_lens_cpu.max())
+
+        self.workspace = torch.empty(
+            384 * 1024 * 1024, dtype=torch.uint8, device=self.device
+        )
+
+        # Shared random q/kv so both code paths attend over identical data.
+        self.q = torch.randn(
+            self.total_q, NUM_QO_HEADS, HEAD_DIM, dtype=DTYPE, device=self.device
+        )
+        # page_size == 1 -> [num_pages, 1, num_kv_heads, head_dim] (NHD).
+        self.k_cache = torch.randn(
+            self.total_kv, 1, NUM_KV_HEADS, HEAD_DIM, dtype=DTYPE, device=self.device
+        )
+        self.v_cache = torch.randn(
+            self.total_kv, 1, NUM_KV_HEADS, HEAD_DIM, dtype=DTYPE, device=self.device
+        )
+
+    def _new_wrapper(self):
+        bs = self.bs
+        return BatchPrefillWithPagedKVCacheWrapper(
+            self.workspace,
+            "NHD",
+            use_cuda_graph=True,
+            backend="fa2",
+            qo_indptr_buf=torch.zeros(bs + 1, dtype=torch.int32, device=self.device),
+            paged_kv_indptr_buf=torch.zeros(
+                bs + 1, dtype=torch.int32, device=self.device
+            ),
+            paged_kv_indices_buf=torch.zeros(
+                self.total_kv, dtype=torch.int32, device=self.device
+            ),
+            paged_kv_last_page_len_buf=torch.ones(
+                bs, dtype=torch.int32, device=self.device
+            ),
+        )
+
+    def _real_plan(self, w):
+        w.plan(
+            self.qo_indptr,
+            self.kv_indptr,
+            self.kv_indices,
+            self.last_page_len,
+            NUM_QO_HEADS,
+            NUM_KV_HEADS,
+            HEAD_DIM,
+            1,  # page_size
+            causal=True,
+            q_data_type=DTYPE,
+            kv_data_type=DTYPE,
+        )
+
+    def _forward(self, w):
+        return w.run(self.q, (self.k_cache, self.v_cache))
+
+    def _out_upstream(self):
+        """Ground truth: a wrapper planned only by upstream plan()."""
+        w = self._new_wrapper()
+        self._real_plan(w)
+        return self._forward(w)
+
+    def _out_fast(self, *, kv_indices=None):
+        """Same attention, planned via the host-known fast path. One real plan()
+        first populates `_cached_module` (mirrors capture), then fast_prefill_plan
+        re-plans from host metadata."""
+        if kv_indices is None:
+            kv_indices = self.kv_indices
+        w = self._new_wrapper()
+        self._real_plan(w)
+        fast_prefill_plan(
+            w,
+            self.qo_indptr,
+            self.kv_indptr,
+            kv_indices,
+            self.last_page_len,
+            NUM_QO_HEADS,
+            NUM_KV_HEADS,
+            HEAD_DIM,
+            1,
+            causal=True,
+            q_data_type=DTYPE,
+            kv_data_type=DTYPE,
+            qo_indptr_host=self.qo_indptr_host,
+            kv_indptr_host=self.kv_indptr_host,
+            kv_lens_host=self.kv_lens_host,
+            max_q_len=self.max_q_len,
+            max_kv_len=self.max_kv_len,
+        )
+        return self._forward(w)
+
+    def test_fast_plan_matches_upstream(self):
+        # Two genuinely independent plan paths over identical q/kv must produce
+        # the same attention output.
+        out_upstream = self._out_upstream()
+        out_fast = self._out_fast()
+        torch.testing.assert_close(out_fast, out_upstream, rtol=0, atol=0)
+
+    def test_mutation_changes_output(self):
+        """Guards against a vacuous test: the kv_indices fast_prefill_plan installs
+        select which physical KV slots the kernel gathers, so reversing them must
+        change the attention output. If it did not, the equivalence assertion
+        would not be exercising the metadata fast_prefill_plan is responsible for."""
+        out_upstream = self._out_upstream()
+        reversed_kv = torch.flip(self.kv_indices, dims=[0]).contiguous()
+        out_wrong = self._out_fast(kv_indices=reversed_kv)
+        self.assertFalse(
+            torch.allclose(out_wrong, out_upstream, rtol=1e-3, atol=1e-3),
+            "output unchanged under reversed kv_indices; test lacks discriminating power",
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()