[Spec] Add sync-free fast_prefill_plan for EAGLE draft-extend CUDA graph (#28854)
This commit is contained in:
Liangsheng Yin
@@ -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,
|
||||
)
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user