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

python/sglang/srt/layers/attention/flashinfer_backend.py

@@ -175,6 +175,119 @@ global_workspace_buffer = None
 global_override_indptr_cpu = None
 
 
+def fast_prefill_plan(
+    self,
+    qo_indptr: torch.Tensor,
+    paged_kv_indptr: torch.Tensor,
+    paged_kv_indices: torch.Tensor,
+    paged_kv_last_page_len: torch.Tensor,
+    num_qo_heads: int,
+    num_kv_heads: int,
+    head_dim_qk: int,
+    page_size: int,
+    head_dim_vo: Optional[int] = None,
+    custom_mask: Optional[torch.Tensor] = None,
+    causal: bool = False,
+    window_left: int = -1,
+    q_data_type: Union[str, torch.dtype] = "float16",
+    kv_data_type: Optional[Union[str, torch.dtype]] = None,
+    o_data_type: Optional[Union[str, torch.dtype]] = None,
+    non_blocking: bool = True,
+    fixed_split_size: Optional[int] = None,
+    prefix_len_ptr: Optional[torch.Tensor] = None,
+    token_pos_in_items_ptr: Optional[torch.Tensor] = None,
+    token_pos_in_items_len: int = 0,
+    max_item_len_ptr: Optional[torch.Tensor] = None,
+    # Required host-known metadata: lets us skip the per-replay device-to-host
+    # copies upstream plan() always issues. Keyword-only with no default so a
+    # caller that forgets them fails at the call boundary, not with a cryptic
+    # None crash deeper in.
+    *,
+    qo_indptr_host: torch.Tensor,
+    kv_indptr_host: torch.Tensor,
+    kv_lens_host: torch.Tensor,
+    max_q_len: int,
+    max_kv_len: int,
+) -> None:
+    """Sync-free ``BatchPrefillWithPagedKVCacheWrapper.plan`` for the EAGLE
+    draft-extend CUDA graph (FlashInfer fa2, cuda-graph mode only).
+
+    Upstream plan() always does qo/paged_kv/last_page_len ``.to("cpu")`` to build
+    its host scheduling metadata, a blocking D2H that drains the GPU queue every
+    replay. The caller passes host-known qo/kv layout in, so we call the underlying
+    ``_cached_module.plan`` directly with no readback; the ``_plan_info`` produced
+    is identical to plan()'s.
+    """
+    assert self.is_cuda_graph_enabled, "fast_prefill_plan is cuda-graph only"
+    assert (
+        getattr(self, "_backend", None) == "fa2"
+    ), "fast_prefill_plan supports the fa2 backend only"
+    assert (
+        getattr(self, "_cached_module", None) is not None
+    ), "fast_prefill_plan requires _cached_module from a prior real plan() (capture)"
+
+    if head_dim_vo is None:
+        head_dim_vo = head_dim_qk
+    batch_size = len(paged_kv_last_page_len)
+
+    total_num_rows = int(qo_indptr_host[-1])
+    self._qo_indptr_last = total_num_rows
+    self._max_q_len = max_q_len
+    self._max_kv_len = max_kv_len
+
+    if self._max_total_num_rows is None:
+        self._max_total_num_rows = total_num_rows
+
+    self._batch_size = batch_size
+    self._num_qo_heads = num_qo_heads
+    self._num_kv_heads = num_kv_heads
+    self._prefix_len_ptr = prefix_len_ptr
+    self._token_pos_in_items_ptr = token_pos_in_items_ptr
+    self._token_pos_in_items_len = token_pos_in_items_len
+    self._max_item_len_ptr = max_item_len_ptr
+
+    # Refresh the cuda-graph input buffers (device-to-device, non-blocking).
+    self._qo_indptr_buf.copy_(qo_indptr, non_blocking=non_blocking)
+    self._paged_kv_indptr_buf.copy_(paged_kv_indptr, non_blocking=non_blocking)
+    self._paged_kv_last_page_len_buf.copy_(
+        paged_kv_last_page_len, non_blocking=non_blocking
+    )
+    self._paged_kv_indices_buf[: len(paged_kv_indices)].copy_(
+        paged_kv_indices,
+        non_blocking=(paged_kv_indices.device == self.device) and non_blocking,
+    )
+
+    self._cached_q_data_type = q_data_type
+    self._cached_kv_data_type = (
+        kv_data_type if kv_data_type is not None else q_data_type
+    )
+    self._cached_o_data_type = o_data_type
+    self._block_tables = None
+
+    args = [
+        self._float_workspace_buffer,
+        self._int_workspace_buffer,
+        self._pin_memory_int_workspace_buffer,
+        qo_indptr_host,
+        kv_indptr_host,
+        kv_lens_host,
+        self._max_total_num_rows or total_num_rows,
+        batch_size,
+        num_qo_heads,
+        num_kv_heads,
+        page_size,
+        self.is_cuda_graph_enabled,
+        head_dim_qk,
+        head_dim_vo,
+        causal,
+        window_left,
+        fixed_split_size if fixed_split_size is not None else -1,
+        False,  # disable_split_kv
+        0,  # num_colocated_ctas
+    ]
+    self._plan_info = self._cached_module.plan(*args)
+
+
 class FlashInferAttnBackend(AttentionBackend):
     """Flashinfer attention kernels."""
 
@@ -593,6 +706,20 @@ class FlashInferAttnBackend(AttentionBackend):
             for w in self.decode_cuda_graph_metadata[bs]:
                 w.begin_forward = partial(fast_decode_plan, w)
 
+        if (
+            in_capture
+            and forward_mode.is_draft_extend_v2()
+            and self.prefill_backend == "fa2"
+            # Host-rebuilt layout only matches full attention (single wrapper);
+            # SWA/cross-attn keep the plain plan().
+            and self.dispatch_reason is None
+        ):
+            # Like decode: swap in fast_prefill_plan for replay, after the real
+            # plan() above set up _cached_module (host metadata supplied per-replay
+            # in call_begin_forward).
+            for w in self.draft_extend_cuda_graph_metadata[bs]:
+                w.begin_forward = partial(fast_prefill_plan, w)
+
         # Refill the SWA write-target buffer from the live out_cache_loc before
         # replay (bound onto the metadata at capture below).
         if self.use_sliding_window_kv_pool and forward_batch.out_cache_loc is not None:
@@ -1372,6 +1499,7 @@ class FlashInferIndicesUpdaterPrefill:
             spec_info,
             fixed_split_size=fixed_split_size,
             multi_item_params=multi_item_params,
+            seq_lens_cpu=seq_lens_cpu,
         )
 
     def update_sliding_window(
@@ -1561,6 +1689,7 @@ class FlashInferIndicesUpdaterPrefill:
         fixed_split_size: Optional[int] = None,
         multi_item_params: Optional[MultiItemScoringParams] = None,
         cross_attention_custom_mask: Optional[torch.Tensor] = None,
+        seq_lens_cpu: Optional[torch.Tensor] = None,
     ):
         bs = len(seq_lens)
         if spec_info is None:
@@ -1646,6 +1775,40 @@ class FlashInferIndicesUpdaterPrefill:
             token_pos_in_items_len = 0
             max_item_len_ptr = None
 
+        # fast_prefill_plan (installed at capture) is sync-free: it needs the
+        # host-known qo/kv layout from the caller. Assert rather than silently
+        # fall back to plan()'s blocking D2H on the replay hot-path.
+        paged_plan_kwargs = {}
+        num_tokens_per_req = getattr(spec_info, "num_tokens_per_req", None)
+        uses_fast_prefill = (
+            hasattr(wrapper_paged.begin_forward, "func")
+            and wrapper_paged.begin_forward.func is fast_prefill_plan
+        )
+        if uses_fast_prefill:
+            assert (
+                seq_lens_cpu is not None
+            ), "fast_prefill_plan replay requires host-known seq_lens_cpu (got None)"
+            assert (
+                num_tokens_per_req is not None and num_tokens_per_req > 0
+            ), f"fast_prefill_plan replay requires num_tokens_per_req > 0 (got {num_tokens_per_req})"
+            seq_lens_cpu_i32 = seq_lens_cpu.to(torch.int32)
+            qo_indptr_host = torch.arange(
+                0,
+                (bs + 1) * num_tokens_per_req,
+                step=num_tokens_per_req,
+                dtype=torch.int32,
+                device="cpu",
+            )
+            kv_indptr_host = torch.zeros(bs + 1, dtype=torch.int32, device="cpu")
+            kv_indptr_host[1:] = torch.cumsum(seq_lens_cpu_i32, dim=0)
+            paged_plan_kwargs = dict(
+                qo_indptr_host=qo_indptr_host,
+                kv_indptr_host=kv_indptr_host,
+                kv_lens_host=seq_lens_cpu_i32,
+                max_q_len=num_tokens_per_req,
+                max_kv_len=int(seq_lens_cpu_i32.max()),
+            )
+
         wrapper_paged.begin_forward(
             qo_indptr,
             kv_indptr,
@@ -1664,6 +1827,7 @@ class FlashInferIndicesUpdaterPrefill:
             token_pos_in_items_ptr=token_pos_in_items_ptr,
             token_pos_in_items_len=token_pos_in_items_len,
             max_item_len_ptr=max_item_len_ptr,
+            **paged_plan_kwargs,
         )
 
 

test/registered/unit/spec/test_fast_prefill_plan.py

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