minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Liangsheng Yin
2026-06-22 15:15:15 -07:00
committed by GitHub
parent c0198fc277
commit 770d6b2825
2 changed files with 351 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/flashinfer_backend.py
+164
View File
@@ -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
+187
View File
@@ -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()