188 lines
6.8 KiB
Python
188 lines
6.8 KiB
Python
"""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=10, 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()
|