diff --git a/python/sglang/srt/layers/attention/flashinfer_backend.py b/python/sglang/srt/layers/attention/flashinfer_backend.py index e8a919a84..da22da628 100644 --- a/python/sglang/srt/layers/attention/flashinfer_backend.py +++ b/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, ) diff --git a/test/registered/unit/spec/test_fast_prefill_plan.py b/test/registered/unit/spec/test_fast_prefill_plan.py new file mode 100644 index 000000000..9b3c14267 --- /dev/null +++ b/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()