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

[refactor] unify cuda-graph capture/replay across attention backends (#26134)

Browse Source 
Co-authored-by: Cheng Wan <cheng.wan@radixark.ai>
Co-authored-by: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Cheng Wan
2026-05-22 18:51:13 -07:00
committed by GitHub
co-authored by Cheng Wan Claude Sonnet 4.6
parent 208397affc
commit d226f75669
5 changed files with 461 additions and 615 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/cutlass_mla_backend.py
+17 -23
View File
@@ -153,24 +153,22 @@ class CutlassMLABackend(FlashInferMLAAttnBackend):
forward_mode: ForwardMode,
spec_info: Optional[SpecInput],
):
if forward_mode.is_decode_or_idle():
if spec_info is None:
max_seqlen_pad = self.cuda_graph_kv_indices.shape[1]
create_flashmla_kv_indices_triton[(bs,)](
self.req_to_token,
req_pool_indices,
seq_lens,
None,
self.cuda_graph_kv_indices,
self.req_to_token.stride(0),
self.cuda_graph_kv_indices.stride(0),
PAGED_SIZE=PAGE_SIZE,
)
self.forward_metadata = CutlassMLADecodeMetadata(
self.cuda_graph_mla_workspace,
self.cuda_graph_kv_indices[:bs, :max_seqlen_pad],
)
if forward_mode.is_decode_or_idle() and spec_info is None:
self.init_forward_metadata_replay_cuda_graph(
bs=bs,
req_pool_indices=req_pool_indices,
seq_lens=seq_lens,
seq_lens_sum=None,
encoder_lens=encoder_lens,
forward_mode=forward_mode,
spec_info=spec_info,
seq_lens_cpu=None,
)
max_seqlen_pad = self.cuda_graph_kv_indices.shape[1]
self.forward_metadata = CutlassMLADecodeMetadata(
self.cuda_graph_mla_workspace,
self.cuda_graph_kv_indices[:bs, :max_seqlen_pad],
)
else:
super().init_forward_metadata_capture_cuda_graph(
bs,
@@ -193,15 +191,11 @@ class CutlassMLABackend(FlashInferMLAAttnBackend):
spec_info: Optional[SpecInput],
seq_lens_cpu: Optional[torch.Tensor],
):
if forward_mode.is_decode_or_idle():
assert seq_lens_cpu is not None
seq_lens = seq_lens[:bs]
create_flashmla_kv_indices_triton[(bs,)](
self.req_to_token,
req_pool_indices[:bs],
seq_lens,
seq_lens[:bs],
None,
self.cuda_graph_kv_indices,
self.req_to_token.stride(0),
python/sglang/srt/layers/attention/flashinfer_backend.py
+93 -154
View File
@@ -559,6 +559,81 @@ class FlashInferAttnBackend(AttentionBackend):
self.cuda_graph_qk_indptr = [x.clone() for x in self.kv_indptr]
self.cuda_graph_qo_indptr = [x.clone() for x in self.kv_indptr]
def _create_decode_wrappers(self, bs: int, num_tokens: int) -> list:
return [
BatchDecodeWithPagedKVCacheWrapper(
self.workspace_buffer,
"NHD",
backend=self.decode_backend,
use_cuda_graph=True,
use_tensor_cores=self.decode_use_tensor_cores,
paged_kv_indptr_buffer=self.kv_indptr[i][: num_tokens + 1],
paged_kv_indices_buffer=self.cuda_graph_kv_indices[i],
paged_kv_last_page_len_buffer=self.kv_last_page_len[:num_tokens],
)
for i in range(self.num_wrappers)
]
def _create_prefill_wrappers(self, bs: int, use_custom_mask: bool = False) -> list:
# FlashInfer's prefill wrapper decides mask mode based on whether
# `custom_mask_buf` is initialized (not whether a custom mask is provided).
# For cases like DFLASH draft (ENCODER_ONLY / non-causal) we do NOT use a
# custom mask, so we must avoid initializing `custom_mask_buf`, otherwise
# FlashInfer will treat the (zero) buffer as a real mask and block attention.
wrappers = []
for i in range(self.num_wrappers):
extra = (
{
"custom_mask_buf": self.cuda_graph_custom_mask,
"mask_indptr_buf": self.cuda_graph_qk_indptr[i][: bs + 1],
}
if use_custom_mask
else {}
)
wrappers.append(
BatchPrefillWithPagedKVCacheWrapper(
self.workspace_buffer,
"NHD",
use_cuda_graph=True,
backend=self.prefill_backend,
qo_indptr_buf=self.cuda_graph_qo_indptr[i][: bs + 1],
paged_kv_indptr_buf=self.kv_indptr[i][: bs + 1],
paged_kv_indices_buf=self.cuda_graph_kv_indices[i],
paged_kv_last_page_len_buf=self.kv_last_page_len[:bs],
**extra,
)
)
return wrappers
def _prepare_cuda_graph_metadata(
self,
bs: int,
num_tokens: int,
forward_mode: ForwardMode,
spec_info: Optional[SpecInput],
) -> None:
if forward_mode.is_decode_or_idle():
decode_wrappers = self._create_decode_wrappers(bs, num_tokens)
self.decode_cuda_graph_metadata[bs] = decode_wrappers
self.forward_metadata = DecodeMetadata(decode_wrappers)
elif (
forward_mode.is_target_verify()
or forward_mode.is_draft_extend()
or forward_mode.is_dllm_extend()
):
use_custom_mask = (
forward_mode.is_target_verify()
and spec_info is not None
and getattr(spec_info, "custom_mask", None) is not None
)
prefill_wrappers = self._create_prefill_wrappers(bs, use_custom_mask)
self.prefill_cuda_graph_metadata[bs] = prefill_wrappers
self.forward_metadata = PrefillMetadata(
prefill_wrappers, forward_mode.is_dllm_extend(), False
)
else:
raise ValueError(f"Invalid mode: {forward_mode=}")
def init_forward_metadata_capture_cuda_graph(
self,
bs: int,
@@ -569,148 +644,24 @@ class FlashInferAttnBackend(AttentionBackend):
forward_mode: ForwardMode,
spec_info: Optional[SpecInput],
):
seq_lens_sum = seq_lens.sum().item()
seq_lens_cpu = seq_lens.cpu()
self._prepare_cuda_graph_metadata(bs, num_tokens, forward_mode, spec_info)
self.init_forward_metadata_replay_cuda_graph(
bs=bs,
req_pool_indices=req_pool_indices,
seq_lens=seq_lens,
seq_lens_sum=seq_lens_sum,
encoder_lens=encoder_lens,
forward_mode=forward_mode,
spec_info=spec_info,
seq_lens_cpu=seq_lens_cpu,
)
# fast_decode_plan requires _cached_module set by the initial full
# begin_forward call above; install it only after that first plan runs.
if forward_mode.is_decode_or_idle():
decode_wrappers = []
for i in range(self.num_wrappers):
decode_wrappers.append(
BatchDecodeWithPagedKVCacheWrapper(
self.workspace_buffer,
"NHD",
backend=self.decode_backend,
use_cuda_graph=True,
use_tensor_cores=self.decode_use_tensor_cores,
paged_kv_indptr_buffer=self.kv_indptr[i][: num_tokens + 1],
paged_kv_indices_buffer=self.cuda_graph_kv_indices[i],
paged_kv_last_page_len_buffer=self.kv_last_page_len[
:num_tokens
],
)
)
seq_lens_sum = seq_lens.sum().item()
self.indices_updater_decode.update(
req_pool_indices,
seq_lens,
seq_lens.cpu(), # may add a little overhead in capture stage
seq_lens_sum,
decode_wrappers=decode_wrappers,
encoder_lens=encoder_lens,
spec_info=spec_info,
fixed_split_size=None,
disable_split_kv=self.disable_cuda_graph_kv_split,
)
self.decode_cuda_graph_metadata[bs] = decode_wrappers
self.forward_metadata = DecodeMetadata(decode_wrappers)
for i in range(self.num_wrappers):
decode_wrappers[i].begin_forward = partial(
fast_decode_plan, decode_wrappers[i]
)
elif forward_mode.is_target_verify():
# FlashInfer's prefill wrapper decides mask mode based on whether
# `custom_mask_buf` is initialized (not whether a custom mask is provided).
# For cases like DFLASH draft (ENCODER_ONLY / non-causal) we do NOT use a
# custom mask, so we must avoid initializing `custom_mask_buf`, otherwise
# FlashInfer will treat the (zero) buffer as a real mask and block attention.
use_custom_mask = (
spec_info is not None
and getattr(spec_info, "custom_mask", None) is not None
)
prefill_wrappers = []
for i in range(self.num_wrappers):
wrapper_kwargs = {}
if use_custom_mask:
wrapper_kwargs = {
"custom_mask_buf": self.cuda_graph_custom_mask,
"mask_indptr_buf": self.cuda_graph_qk_indptr[i][: bs + 1],
}
prefill_wrappers.append(
BatchPrefillWithPagedKVCacheWrapper(
self.workspace_buffer,
"NHD",
use_cuda_graph=True,
backend=self.prefill_backend,
qo_indptr_buf=self.cuda_graph_qo_indptr[i][: bs + 1],
paged_kv_indptr_buf=self.kv_indptr[i][: bs + 1],
paged_kv_indices_buf=self.cuda_graph_kv_indices[i],
paged_kv_last_page_len_buf=self.kv_last_page_len[:bs],
**wrapper_kwargs,
)
)
seq_lens_sum = seq_lens.sum().item()
self.indices_updater_prefill.update(
req_pool_indices,
seq_lens,
seq_lens.cpu(), # may add a little overhead in capture stage
seq_lens_sum,
prefix_lens=None,
prefill_wrappers=prefill_wrappers,
use_ragged=False,
encoder_lens=encoder_lens,
spec_info=spec_info,
)
self.prefill_cuda_graph_metadata[bs] = prefill_wrappers
self.forward_metadata = PrefillMetadata(prefill_wrappers, False, False)
elif forward_mode.is_draft_extend():
prefill_wrappers = []
for i in range(self.num_wrappers):
prefill_wrappers.append(
BatchPrefillWithPagedKVCacheWrapper(
self.workspace_buffer,
"NHD",
backend=self.prefill_backend,
use_cuda_graph=True,
qo_indptr_buf=self.cuda_graph_qo_indptr[i][: bs + 1],
paged_kv_indptr_buf=self.kv_indptr[i][: bs + 1],
paged_kv_indices_buf=self.cuda_graph_kv_indices[i],
paged_kv_last_page_len_buf=self.kv_last_page_len[:bs],
)
)
seq_lens_sum = seq_lens.sum().item()
self.indices_updater_prefill.update(
req_pool_indices,
seq_lens,
seq_lens.cpu(), # may add a little overhead in capture stage
seq_lens_sum,
prefix_lens=None,
prefill_wrappers=prefill_wrappers,
use_ragged=False,
encoder_lens=encoder_lens,
spec_info=spec_info,
)
self.prefill_cuda_graph_metadata[bs] = prefill_wrappers
self.forward_metadata = PrefillMetadata(prefill_wrappers, False, False)
elif forward_mode.is_dllm_extend():
prefill_wrappers = []
for i in range(self.num_wrappers):
prefill_wrappers.append(
BatchPrefillWithPagedKVCacheWrapper(
self.workspace_buffer,
"NHD",
backend=self.prefill_backend,
use_cuda_graph=True,
qo_indptr_buf=self.cuda_graph_qo_indptr[i][: bs + 1],
paged_kv_indptr_buf=self.kv_indptr[i][: bs + 1],
paged_kv_indices_buf=self.cuda_graph_kv_indices[i],
paged_kv_last_page_len_buf=self.kv_last_page_len[:bs],
)
)
seq_lens_sum = seq_lens.sum().item()
self.indices_updater_prefill.update(
req_pool_indices,
seq_lens,
seq_lens.cpu(), # may add a little overhead in capture stage
seq_lens_sum,
prefix_lens=seq_lens - self.dllm_config.block_size,
prefill_wrappers=prefill_wrappers,
use_ragged=not self.use_paged,
encoder_lens=encoder_lens,
spec_info=None,
)
self.prefill_cuda_graph_metadata[bs] = prefill_wrappers
self.forward_metadata = PrefillMetadata(prefill_wrappers, True, False)
else:
raise ValueError(f"Invalid mode: {forward_mode=}")
for w in self.decode_cuda_graph_metadata[bs]:
w.begin_forward = partial(fast_decode_plan, w)
def init_forward_metadata_replay_cuda_graph(
self,
@@ -735,19 +686,7 @@ class FlashInferAttnBackend(AttentionBackend):
fixed_split_size=None,
disable_split_kv=self.disable_cuda_graph_kv_split,
)
elif forward_mode.is_target_verify():
self.indices_updater_prefill.update(
req_pool_indices[:bs],
seq_lens[:bs],
seq_lens_cpu[:bs] if seq_lens_cpu is not None else None,
seq_lens_sum,
prefix_lens=None,
prefill_wrappers=self.prefill_cuda_graph_metadata[bs],
use_ragged=False,
encoder_lens=encoder_lens[:bs] if encoder_lens is not None else None,
spec_info=spec_info,
)
elif forward_mode.is_draft_extend():
elif forward_mode.is_target_verify() or forward_mode.is_draft_extend():
self.indices_updater_prefill.update(
req_pool_indices[:bs],
seq_lens[:bs],
python/sglang/srt/layers/attention/flashinfer_mla_backend.py
+6 -38
View File
@@ -407,8 +407,8 @@ class FlashInferMLAAttnBackend(AttentionBackend):
self.decode_cuda_graph_metadata[bs] = decode_wrapper
self.forward_metadata = DecodeMetadata(decode_wrapper)
decode_wrapper.plan = partial(fast_mla_decode_plan, decode_wrapper)
elif forward_mode.is_target_verify():
verify_wrapper = BatchMLAPagedAttentionWrapper(
elif forward_mode.is_target_verify() or forward_mode.is_draft_extend():
prefill_wrapper = BatchMLAPagedAttentionWrapper(
self.workspace_buffer,
use_cuda_graph=True,
qo_indptr=self.cuda_graph_qo_indptr[: bs + 1],
@@ -423,34 +423,12 @@ class FlashInferMLAAttnBackend(AttentionBackend):
seq_lens,
seq_lens_sum,
prefix_lens=None,
prefill_wrapper_paged=verify_wrapper,
prefill_wrapper_paged=prefill_wrapper,
use_ragged=False,
spec_info=spec_info,
)
self.prefill_cuda_graph_metadata[bs] = verify_wrapper
self.forward_metadata = PrefillMetadata(verify_wrapper, False)
elif forward_mode.is_draft_extend():
draft_extend_wrapper = BatchMLAPagedAttentionWrapper(
self.workspace_buffer,
use_cuda_graph=True,
qo_indptr=self.cuda_graph_qo_indptr[: bs + 1],
kv_indptr=self.cuda_graph_kv_indptr[: bs + 1],
kv_indices=self.cuda_graph_kv_indices,
kv_len_arr=self.cuda_graph_kv_lens[:bs],
backend="auto",
)
seq_lens_sum = seq_lens.sum().item()
self.indices_updater_prefill.update(
req_pool_indices,
seq_lens,
seq_lens_sum,
prefix_lens=None,
prefill_wrapper_paged=draft_extend_wrapper,
use_ragged=False,
spec_info=spec_info,
)
self.prefill_cuda_graph_metadata[bs] = draft_extend_wrapper
self.forward_metadata = PrefillMetadata(draft_extend_wrapper, False)
self.prefill_cuda_graph_metadata[bs] = prefill_wrapper
self.forward_metadata = PrefillMetadata(prefill_wrapper, False)
else:
raise ValueError(f"Invalid mode: {forward_mode=}")
@@ -488,17 +466,7 @@ class FlashInferMLAAttnBackend(AttentionBackend):
spec_info=spec_info,
**self.fast_decode_kwargs,
)
elif forward_mode.is_target_verify():
self.indices_updater_prefill.update(
req_pool_indices[:bs],
seq_lens[:bs],
seq_lens_sum,
prefix_lens=None,
prefill_wrapper_paged=self.prefill_cuda_graph_metadata[bs],
use_ragged=False,
spec_info=spec_info,
)
elif forward_mode.is_draft_extend():
elif forward_mode.is_target_verify() or forward_mode.is_draft_extend():
self.indices_updater_prefill.update(
req_pool_indices[:bs],
seq_lens[:bs],
python/sglang/srt/layers/attention/triton_backend.py
+286 -329
View File
@@ -283,11 +283,152 @@ class TritonAttnBackend(AttentionBackend):
MAX_NUM_SEQ=SCHEDULE_SEQ,
)
def _fill_kv_indptr_and_indices(
self,
bs: int,
seq_lens: torch.Tensor,
req_pool_indices: torch.Tensor,
kv_indices: torch.Tensor,
) -> torch.Tensor:
kv_indptr = self.kv_indptr[: bs + 1]
kv_indptr[1:] = torch.cumsum(seq_lens, dim=0)
create_flashinfer_kv_indices_triton[(bs,)](
self.req_to_token,
req_pool_indices,
seq_lens,
kv_indptr,
None,
kv_indices,
self.req_to_token.stride(0),
)
return kv_indptr
def _update_decode_kv_buffers(
self,
bs: int,
seq_lens: torch.Tensor,
req_pool_indices: torch.Tensor,
):
"""Fill KV (and SWA) cuda-graph buffers for decode/idle mode.
Returns ``(kv_indptr, window_kv_indptr, window_kv_lens)`` where
``window_kv_lens`` is ``None`` when sliding-window is disabled.
"""
seq_lens = seq_lens[:bs]
req_pool_indices = req_pool_indices[:bs]
kv_indptr = self._fill_kv_indptr_and_indices(
bs, seq_lens, req_pool_indices, self.cuda_graph_kv_indices
)
window_kv_indptr = self.window_kv_indptr
window_kv_lens = None
if self.sliding_window_size is not None and self.sliding_window_size > 0:
window_kv_indptr, _, window_kv_lens, _ = update_sliding_window_buffer(
self.window_kv_indptr,
self.req_to_token,
self.sliding_window_size,
seq_lens,
req_pool_indices,
bs,
token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
window_kv_indices=self.cuda_graph_window_kv_indices,
)
return kv_indptr, window_kv_indptr, window_kv_lens
def _update_target_verify_buffers(
self,
bs: int,
seq_lens: torch.Tensor,
req_pool_indices: torch.Tensor,
spec_info,
):
"""Fill all cuda-graph buffers for target_verify mode.
Returns the ForwardMetadata components:
``(qo_indptr, kv_indptr, custom_mask, mask_indptr,
window_kv_indptr, window_kv_indices, window_num_kv_splits, window_kv_offsets)``
"""
qo_indptr = self.qo_indptr[: bs + 1]
qo_indptr[: bs + 1] = torch.arange(
0,
(1 + bs) * self.num_draft_tokens,
step=self.num_draft_tokens,
dtype=torch.int32,
device=self.device,
)
kv_indptr = self._fill_kv_indptr_and_indices(
bs, seq_lens, req_pool_indices, self.cuda_graph_kv_indices
)
window_kv_indptr = self.window_kv_indptr
window_kv_indices = None
window_num_kv_splits = None
window_kv_offsets = None
if self.sliding_window_size is not None and self.sliding_window_size > 0:
window_kv_indices = self.cuda_graph_window_kv_indices
window_num_kv_splits = self.cuda_graph_window_num_kv_splits
window_kv_offsets = self.cuda_graph_window_kv_offsets
window_kv_indptr, window_kv_indices, _, window_kv_offsets[:bs] = (
update_sliding_window_buffer(
self.window_kv_indptr,
self.req_to_token,
self.sliding_window_size,
seq_lens[:bs],
req_pool_indices,
bs,
token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
window_kv_indices=window_kv_indices,
)
)
custom_mask = self.cuda_graph_custom_mask
if (
spec_info is not None
and getattr(spec_info, "custom_mask", None) is not None
):
custom_mask[: spec_info.custom_mask.shape[0]] = spec_info.custom_mask
else:
custom_mask = None
seq_mask_len = self.num_draft_tokens * (seq_lens + self.num_draft_tokens)
mask_indptr = self.mask_indptr[: bs + 1]
mask_indptr[1 : bs + 1] = torch.cumsum(seq_mask_len, dim=0)
return (
qo_indptr,
kv_indptr,
custom_mask,
mask_indptr,
window_kv_indptr,
window_kv_indices,
window_num_kv_splits,
window_kv_offsets,
)
def _update_draft_extend_buffers(
self,
bs: int,
seq_lens: torch.Tensor,
req_pool_indices: torch.Tensor,
):
"""Fill QO + KV cuda-graph buffers for draft_extend mode.
Returns ``(qo_indptr, kv_indptr, num_tokens_per_bs)``.
"""
seq_lens = seq_lens[:bs]
num_tokens_per_bs = self.speculative_num_steps + 1
qo_indptr = self.qo_indptr[: bs + 1]
qo_indptr[: bs + 1] = torch.arange(
0,
bs * num_tokens_per_bs + 1,
step=num_tokens_per_bs,
dtype=torch.int32,
device=self.device,
)
kv_indptr = self._fill_kv_indptr_and_indices(
bs, seq_lens, req_pool_indices, self.cuda_graph_kv_indices
)
return qo_indptr, kv_indptr, num_tokens_per_bs
def init_forward_metadata(self, forward_batch: ForwardBatch):
"""Init auxiliary variables for triton attention backend."""
bs = forward_batch.batch_size
kv_indptr = self.kv_indptr
window_kv_indptr = self.window_kv_indptr
window_kv_indices = None
window_num_kv_splits = None
@@ -297,19 +438,14 @@ class TritonAttnBackend(AttentionBackend):
if forward_batch.forward_mode.is_decode_or_idle():
if spec_info is None:
kv_indptr[1 : bs + 1] = torch.cumsum(forward_batch.seq_lens, dim=0)
kv_indptr = kv_indptr[: bs + 1]
kv_indices = torch.empty(
forward_batch.seq_lens_sum, dtype=torch.int64, device=self.device
)
create_flashinfer_kv_indices_triton[(bs,)](
self.req_to_token,
forward_batch.req_pool_indices,
kv_indptr = self._fill_kv_indptr_and_indices(
bs,
forward_batch.seq_lens,
kv_indptr,
None,
forward_batch.req_pool_indices,
kv_indices,
self.req_to_token.stride(0),
)
# Sliding window
if (
@@ -371,19 +507,14 @@ class TritonAttnBackend(AttentionBackend):
device=self.device,
)
# Different with flashinfer kv_indptr and kv_indices construction
kv_indptr[1 : bs + 1] = torch.cumsum(forward_batch.seq_lens, dim=0)
kv_indptr = kv_indptr[: bs + 1]
kv_indices = torch.empty(
kv_indptr[-1], dtype=torch.int64, device=self.device
forward_batch.seq_lens_sum, dtype=torch.int64, device=self.device
)
create_flashinfer_kv_indices_triton[(bs,)](
self.req_to_token,
forward_batch.req_pool_indices,
kv_indptr = self._fill_kv_indptr_and_indices(
bs,
forward_batch.seq_lens,
kv_indptr,
None,
forward_batch.req_pool_indices,
kv_indices,
self.req_to_token.stride(0),
)
if self.sliding_window_size is not None and self.sliding_window_size > 0:
@@ -435,23 +566,16 @@ class TritonAttnBackend(AttentionBackend):
attn_logits = None
attn_lse = None
else:
kv_indptr[1 : bs + 1] = torch.cumsum(
forward_batch.extend_prefix_lens, dim=0
)
kv_indptr = kv_indptr[: bs + 1]
kv_indices = torch.empty(
sum(forward_batch.extend_prefix_lens_cpu),
dtype=torch.int64,
device=self.device,
)
create_flashinfer_kv_indices_triton[(bs,)](
self.req_to_token,
forward_batch.req_pool_indices,
kv_indptr = self._fill_kv_indptr_and_indices(
bs,
forward_batch.extend_prefix_lens,
kv_indptr,
None,
forward_batch.req_pool_indices,
kv_indices,
self.req_to_token.stride(0),
)
# Sliding window
if self.sliding_window_size is not None and self.sliding_window_size > 0:
@@ -578,6 +702,83 @@ class TritonAttnBackend(AttentionBackend):
device=self.device,
)
def _build_cuda_graph_forward_metadata(
self,
bs: int,
forward_mode: ForwardMode,
spec_info: Optional[SpecInput],
) -> ForwardMetadata:
"""Construct ForwardMetadata from the current cuda-graph buffer state.
Called by capture after the buffer-update helpers have already run
(either via replay or directly). All fields reference the same
``self.cuda_graph_*`` tensors that the captured graph kernels will
read — the Python object is rebuilt each capture, but the underlying
GPU memory addresses are stable.
"""
swa = self.sliding_window_size is not None and self.sliding_window_size > 0
if forward_mode.is_decode_or_idle():
return ForwardMetadata(
attn_logits=self.cuda_graph_attn_logits,
attn_lse=self.cuda_graph_attn_lse,
max_extend_len=None,
num_kv_splits=self.cuda_graph_num_kv_splits,
kv_indptr=self.kv_indptr[: bs + 1],
kv_indices=self.cuda_graph_kv_indices,
qo_indptr=None,
custom_mask=None,
mask_indptr=None,
window_kv_indptr=self.window_kv_indptr[: bs + 1] if swa else None,
window_kv_indices=self.cuda_graph_window_kv_indices if swa else None,
window_num_kv_splits=(
self.cuda_graph_window_num_kv_splits if swa else None
),
window_kv_offsets=None,
swa_attn_logits=self.cuda_graph_swa_attn_logits,
)
elif forward_mode.is_target_verify():
custom_mask = (
self.cuda_graph_custom_mask
if spec_info is not None
and getattr(spec_info, "custom_mask", None) is not None
else None
)
return ForwardMetadata(
attn_logits=None,
attn_lse=None,
max_extend_len=self.num_draft_tokens,
num_kv_splits=None,
kv_indptr=self.kv_indptr[: bs + 1],
kv_indices=self.cuda_graph_kv_indices,
qo_indptr=self.qo_indptr[: bs + 1],
custom_mask=custom_mask,
mask_indptr=self.mask_indptr[: bs + 1],
window_kv_indptr=self.window_kv_indptr[: bs + 1] if swa else None,
window_kv_indices=self.cuda_graph_window_kv_indices if swa else None,
window_num_kv_splits=(
self.cuda_graph_window_num_kv_splits if swa else None
),
window_kv_offsets=self.cuda_graph_window_kv_offsets if swa else None,
)
elif forward_mode.is_draft_extend(include_v2=True):
return ForwardMetadata(
attn_logits=None,
attn_lse=None,
max_extend_len=self.speculative_num_steps + 1,
num_kv_splits=None,
kv_indptr=self.kv_indptr[: bs + 1],
kv_indices=self.cuda_graph_kv_indices,
qo_indptr=self.qo_indptr[: bs + 1],
custom_mask=None,
mask_indptr=None,
window_kv_indptr=self.window_kv_indptr,
window_kv_indices=None,
window_num_kv_splits=None,
window_kv_offsets=None,
)
else:
raise ValueError(f"Invalid forward mode: {forward_mode=} for CUDA Graph.")
def init_forward_metadata_capture_cuda_graph(
self,
bs: int,
@@ -589,158 +790,42 @@ class TritonAttnBackend(AttentionBackend):
spec_info: Optional[SpecInput],
):
assert encoder_lens is None, "Not supported"
window_kv_indptr = self.window_kv_indptr
window_kv_indices = None
window_num_kv_splits = None
window_kv_offsets = None
swa_attn_logits = None
if forward_mode.is_decode_or_idle():
if spec_info is None:
kv_indptr = self.kv_indptr
kv_indptr[1 : bs + 1] = torch.cumsum(seq_lens, dim=0)
kv_indptr = kv_indptr[: bs + 1]
kv_indices = self.cuda_graph_kv_indices
create_flashinfer_kv_indices_triton[(bs,)](
self.req_to_token,
req_pool_indices,
seq_lens,
kv_indptr,
None,
kv_indices,
self.req_to_token.stride(0),
)
if (
self.sliding_window_size is not None
and self.sliding_window_size > 0
):
window_kv_indices = self.cuda_graph_window_kv_indices
window_num_kv_splits = self.cuda_graph_window_num_kv_splits
window_kv_indptr, window_kv_indices, _, _ = (
update_sliding_window_buffer_cuda_graph(
self.window_kv_indptr,
window_kv_indices,
self.req_to_token,
self.sliding_window_size,
seq_lens[:bs],
req_pool_indices,
bs,
self.token_to_kv_pool_allocator,
)
)
else:
kv_indptr, kv_indices = spec_info.kv_indptr, spec_info.kv_indices
# Multi-step speculative decode: kv buffers come from spec_info rather
# than the cuda-graph pool, so replay is not involved for this path.
if forward_mode.is_decode_or_idle() and spec_info is not None:
self.forward_metadata = ForwardMetadata(
attn_logits=self.cuda_graph_attn_logits,
attn_lse=self.cuda_graph_attn_lse,
max_extend_len=None,
num_kv_splits=self.cuda_graph_num_kv_splits,
kv_indptr=spec_info.kv_indptr,
kv_indices=spec_info.kv_indices,
qo_indptr=None,
custom_mask=None,
mask_indptr=None,
window_kv_indptr=self.window_kv_indptr,
window_kv_indices=None,
window_num_kv_splits=None,
window_kv_offsets=None,
swa_attn_logits=self.cuda_graph_swa_attn_logits,
)
return
attn_logits = self.cuda_graph_attn_logits
swa_attn_logits = self.cuda_graph_swa_attn_logits
attn_lse = self.cuda_graph_attn_lse
max_extend_len = None
num_kv_splits = self.cuda_graph_num_kv_splits
qo_indptr = None
custom_mask = None
mask_indptr = None
elif forward_mode.is_target_verify():
qo_indptr = self.qo_indptr[: bs + 1]
qo_indptr[: bs + 1] = torch.arange(
0,
(1 + bs) * self.num_draft_tokens,
step=self.num_draft_tokens,
dtype=torch.int32,
device=self.device,
)
kv_indptr = self.kv_indptr[: bs + 1]
kv_indptr[1 : bs + 1] = torch.cumsum(seq_lens, dim=0)
kv_indices = self.cuda_graph_kv_indices
create_flashinfer_kv_indices_triton[(bs,)](
self.req_to_token,
req_pool_indices,
seq_lens,
kv_indptr,
None,
kv_indices,
self.req_to_token.stride(0),
)
if self.sliding_window_size is not None and self.sliding_window_size > 0:
window_kv_indices = self.cuda_graph_window_kv_indices
window_num_kv_splits = self.cuda_graph_window_num_kv_splits
window_kv_offsets = self.cuda_graph_window_kv_offsets
window_kv_indptr, window_kv_indices, _, window_kv_offsets[:bs] = (
update_sliding_window_buffer_cuda_graph(
self.window_kv_indptr,
window_kv_indices,
self.req_to_token,
self.sliding_window_size,
seq_lens[:bs],
req_pool_indices,
bs,
self.token_to_kv_pool_allocator,
)
)
custom_mask = self.cuda_graph_custom_mask
if (
spec_info is not None
and getattr(spec_info, "custom_mask", None) is not None
):
custom_mask[: spec_info.custom_mask.shape[0]] = spec_info.custom_mask
else:
custom_mask = None
seq_mask_len = self.num_draft_tokens * (seq_lens + self.num_draft_tokens)
mask_indptr = self.mask_indptr[: bs + 1]
mask_indptr[1 : bs + 1] = torch.cumsum(seq_mask_len, dim=0)
max_extend_len = self.num_draft_tokens
num_kv_splits = None
attn_logits = None
attn_lse = None
elif forward_mode.is_draft_extend(include_v2=True):
num_tokens_per_bs = self.speculative_num_steps + 1
qo_indptr = self.qo_indptr[: bs + 1]
qo_indptr[: bs + 1] = torch.arange(
0,
bs * num_tokens_per_bs + 1,
step=num_tokens_per_bs,
dtype=torch.int32,
device=self.device,
)
kv_indptr = self.kv_indptr[: bs + 1]
kv_indptr[1 : bs + 1] = torch.cumsum(seq_lens, dim=0)
kv_indices = self.cuda_graph_kv_indices
create_flashinfer_kv_indices_triton[(bs,)](
self.req_to_token,
req_pool_indices,
seq_lens,
kv_indptr,
None,
kv_indices,
self.req_to_token.stride(0),
)
custom_mask = None
mask_indptr = None
max_extend_len = num_tokens_per_bs
num_kv_splits = None
attn_logits = None
attn_lse = None
else:
raise ValueError(
f"Invalid forward mode: {forward_mode=} for CUDA Graph capture."
)
self.forward_metadata = ForwardMetadata(
attn_logits,
attn_lse,
max_extend_len,
num_kv_splits,
kv_indptr,
kv_indices,
qo_indptr,
custom_mask,
mask_indptr,
window_kv_indptr,
window_kv_indices,
window_num_kv_splits,
window_kv_offsets,
swa_attn_logits=swa_attn_logits,
# Run the same buffer update as replay, then freeze the result into
# a ForwardMetadata whose tensor fields point into the cuda-graph buffers.
self.init_forward_metadata_replay_cuda_graph(
bs=bs,
req_pool_indices=req_pool_indices,
seq_lens=seq_lens,
seq_lens_sum=None,
encoder_lens=encoder_lens,
forward_mode=forward_mode,
spec_info=spec_info,
seq_lens_cpu=None,
)
self.forward_metadata = self._build_cuda_graph_forward_metadata(
bs, forward_mode, spec_info
)
def init_forward_metadata_replay_cuda_graph(
@@ -756,120 +841,22 @@ class TritonAttnBackend(AttentionBackend):
):
# NOTE: encoder_lens expected to be zeros or None
if forward_mode.is_decode_or_idle():
# Update kv_indptr, kv_indices
kv_indptr = self.kv_indptr
kv_indices = self.cuda_graph_kv_indices
num_kv_splits = self.cuda_graph_num_kv_splits
if spec_info is None:
kv_indptr[1 : bs + 1] = torch.cumsum(seq_lens[:bs], dim=0)
kv_indptr = kv_indptr[: bs + 1]
create_flashinfer_kv_indices_triton[(bs,)](
self.req_to_token,
req_pool_indices[:bs],
seq_lens[:bs],
kv_indptr,
None,
kv_indices,
self.req_to_token.stride(0),
assert spec_info is None, "Multi-step cuda graph init is not done here."
_, _, window_kv_lens = self._update_decode_kv_buffers(
bs, seq_lens, req_pool_indices
)
self.get_num_kv_splits(self.cuda_graph_num_kv_splits[:bs], seq_lens[:bs])
if window_kv_lens is not None:
self.get_num_kv_splits(
self.cuda_graph_window_num_kv_splits[:bs], window_kv_lens[:bs]
)
num_token = bs
if (
self.sliding_window_size is not None
and self.sliding_window_size > 0
):
window_num_kv_splits = self.cuda_graph_window_num_kv_splits
window_kv_indices = self.cuda_graph_window_kv_indices
_, _, window_kv_lens, _ = update_sliding_window_buffer_cuda_graph(
self.window_kv_indptr,
window_kv_indices,
self.req_to_token,
self.sliding_window_size,
seq_lens[:bs],
req_pool_indices[:bs],
bs,
self.token_to_kv_pool_allocator,
)
self.get_num_kv_splits(
window_num_kv_splits[:num_token], window_kv_lens[:bs]
)
else:
assert False, "Multi-step cuda graph init is not done here."
self.get_num_kv_splits(num_kv_splits[:num_token], seq_lens[:bs])
elif forward_mode.is_target_verify():
# Update qo_indptr, kv_indptr, kv_indices, custom_mask, mask_indptr
bs = len(req_pool_indices)
qo_indptr = self.qo_indptr[: bs + 1]
qo_indptr[: bs + 1] = torch.arange(
0,
(1 + bs) * self.num_draft_tokens,
step=self.num_draft_tokens,
dtype=torch.int32,
device=self.device,
self._update_target_verify_buffers(
bs, seq_lens, req_pool_indices, spec_info
)
kv_indptr = self.kv_indptr[: bs + 1]
kv_indptr[1 : bs + 1] = torch.cumsum(seq_lens, dim=0)
kv_indices = self.cuda_graph_kv_indices
create_flashinfer_kv_indices_triton[(bs,)](
self.req_to_token,
req_pool_indices,
seq_lens,
kv_indptr,
None,
kv_indices,
self.req_to_token.stride(0),
)
if self.sliding_window_size is not None and self.sliding_window_size > 0:
window_num_kv_splits = self.cuda_graph_window_num_kv_splits
window_kv_indices = self.cuda_graph_window_kv_indices
window_kv_offsets = self.cuda_graph_window_kv_offsets
_, _, window_kv_lens, window_kv_offsets[:bs] = (
update_sliding_window_buffer_cuda_graph(
self.window_kv_indptr,
window_kv_indices,
self.req_to_token,
self.sliding_window_size,
seq_lens[:bs],
req_pool_indices,
bs,
self.token_to_kv_pool_allocator,
)
)
custom_mask = self.cuda_graph_custom_mask
if (
spec_info is not None
and getattr(spec_info, "custom_mask", None) is not None
):
custom_mask[: spec_info.custom_mask.shape[0]] = spec_info.custom_mask
else:
custom_mask = None
seq_mask_len = self.num_draft_tokens * (seq_lens + self.num_draft_tokens)
mask_indptr = self.mask_indptr[: bs + 1]
mask_indptr[1 : bs + 1] = torch.cumsum(seq_mask_len, dim=0)
elif forward_mode.is_draft_extend(include_v2=True):
seq_lens = seq_lens[:bs]
num_tokens_per_bs = self.speculative_num_steps + 1
qo_indptr = self.qo_indptr[: bs + 1]
qo_indptr[: bs + 1] = torch.arange(
0,
bs * num_tokens_per_bs + 1,
step=num_tokens_per_bs,
dtype=torch.int32,
device=self.device,
)
kv_indptr = self.kv_indptr[: bs + 1]
kv_indptr[1 : bs + 1] = torch.cumsum(seq_lens, dim=0)
kv_indices = self.cuda_graph_kv_indices
create_flashinfer_kv_indices_triton[(bs,)](
self.req_to_token,
req_pool_indices,
seq_lens,
kv_indptr,
None,
kv_indices,
self.req_to_token.stride(0),
)
self._update_draft_extend_buffers(bs, seq_lens, req_pool_indices)
else:
raise ValueError(
f"Invalid forward mode: {forward_mode=} for CUDA Graph replay."
@@ -1457,18 +1444,26 @@ def update_sliding_window_buffer(
seq_lens,
req_pool_indices,
bs,
device,
device=None,
token_to_kv_pool_allocator=None,
window_kv_indices=None,
):
"""Fill window KV buffers for sliding-window attention.
Pass ``window_kv_indices`` to write into a pre-allocated buffer (CUDA-graph
path); omit it (or pass ``None``) to allocate a fresh tensor (eager path,
requires ``device``).
"""
window_kv_lens = torch.minimum(
seq_lens,
torch.tensor(sliding_window_size),
)
window_kv_indptr[1 : bs + 1] = torch.cumsum(window_kv_lens, dim=0)
window_kv_indptr = window_kv_indptr[: bs + 1]
window_kv_indices = torch.empty(
window_kv_indptr[-1], dtype=torch.int64, device=device
)
if window_kv_indices is None:
window_kv_indices = torch.empty(
window_kv_indptr[-1], dtype=torch.int64, device=device
)
window_kv_start_idx = seq_lens - window_kv_lens
create_flashinfer_kv_indices_triton[(bs,)](
req_to_token,
@@ -1479,44 +1474,6 @@ def update_sliding_window_buffer(
window_kv_indices,
req_to_token.stride(0),
)
# full to swa index mapping
if hasattr(token_to_kv_pool_allocator, "translate_loc_from_full_to_swa"):
kv_last_index = window_kv_indptr[-1]
window_kv_indices[:kv_last_index] = (
token_to_kv_pool_allocator.translate_loc_from_full_to_swa(
window_kv_indices[:kv_last_index]
)
)
return window_kv_indptr, window_kv_indices, window_kv_lens, window_kv_start_idx
def update_sliding_window_buffer_cuda_graph(
window_kv_indptr,
window_kv_indices,
req_to_token,
sliding_window_size,
seq_lens,
req_pool_indices,
bs,
token_to_kv_pool_allocator=None,
):
window_kv_lens = torch.minimum(
seq_lens,
torch.tensor(sliding_window_size),
)
window_kv_indptr[1 : bs + 1] = torch.cumsum(window_kv_lens, dim=0)
window_kv_indptr = window_kv_indptr[: bs + 1]
window_kv_start_idx = seq_lens - window_kv_lens
create_flashinfer_kv_indices_triton[(bs,)](
req_to_token,
req_pool_indices,
window_kv_lens,
window_kv_indptr,
window_kv_start_idx,
window_kv_indices,
req_to_token.stride(0),
)
# full to swa index mapping
if hasattr(token_to_kv_pool_allocator, "translate_loc_from_full_to_swa"):
kv_last_index = window_kv_indptr[-1]
window_kv_indices[:kv_last_index] = (
python/sglang/srt/layers/attention/wave_backend.py
+59 -71
View File
@@ -390,6 +390,39 @@ class WaveAttnBackend(AttentionBackend):
device=self.device,
)
def _build_cuda_graph_forward_metadata(
self,
bs: int,
forward_mode: ForwardMode,
spec_info: Optional[SpecInput],
) -> ForwardMetadata:
if forward_mode.is_decode_or_idle():
return ForwardMetadata(
attn_logits=self.cuda_graph_attn_logits,
attn_lse=self.cuda_graph_attn_lse,
max_extend_len=None,
num_kv_splits=self.cuda_graph_num_kv_splits,
kv_indptr=self.kv_indptr[: bs + 1],
kv_indices=self.cuda_graph_kv_indices,
qo_indptr=None,
custom_mask=None,
mask_indptr=None,
)
elif forward_mode.is_target_verify():
return ForwardMetadata(
attn_logits=None,
attn_lse=None,
max_extend_len=self.num_draft_tokens,
num_kv_splits=None,
kv_indptr=self.kv_indptr[: bs + 1],
kv_indices=self.cuda_graph_kv_indices,
qo_indptr=self.qo_indptr[: bs + 1],
custom_mask=self.cuda_graph_custom_mask,
mask_indptr=self.mask_indptr[: bs + 1],
)
else:
raise ValueError(f"Invalid forward mode: {forward_mode=} for CUDA Graph.")
def init_forward_metadata_capture_cuda_graph(
self,
bs: int,
@@ -402,76 +435,34 @@ class WaveAttnBackend(AttentionBackend):
):
assert encoder_lens is None, "Not supported"
if forward_mode.is_decode_or_idle():
if spec_info is None:
kv_indptr = self.kv_indptr
kv_indptr[1 : bs + 1] = torch.cumsum(seq_lens, dim=0)
kv_indptr = kv_indptr[: bs + 1]
kv_indices = self.cuda_graph_kv_indices
create_flashinfer_kv_indices_triton[(bs,)](
self.req_to_token,
req_pool_indices,
seq_lens,
kv_indptr,
None,
kv_indices,
self.req_to_token.stride(0),
)
else:
kv_indptr, kv_indices = spec_info.kv_indptr, spec_info.kv_indices
attn_logits = self.cuda_graph_attn_logits
attn_lse = self.cuda_graph_attn_lse
max_extend_len = None
num_kv_splits = self.cuda_graph_num_kv_splits
qo_indptr = None
custom_mask = None
mask_indptr = None
elif forward_mode.is_target_verify():
qo_indptr = self.qo_indptr[: bs + 1]
qo_indptr[: bs + 1] = torch.arange(
0,
(1 + bs) * self.num_draft_tokens,
step=self.num_draft_tokens,
dtype=torch.int32,
device=self.device,
)
kv_indptr = self.kv_indptr[: bs + 1]
kv_indptr[1 : bs + 1] = torch.cumsum(seq_lens, dim=0)
kv_indices = self.cuda_graph_kv_indices
create_flashinfer_kv_indices_triton[(bs,)](
self.req_to_token,
req_pool_indices,
seq_lens,
kv_indptr,
None,
kv_indices,
self.req_to_token.stride(0),
# Multi-step speculative decode: kv buffers come from spec_info rather than
# the cuda-graph pool, so replay is not involved for this path.
if forward_mode.is_decode_or_idle() and spec_info is not None:
self.forward_metadata = ForwardMetadata(
attn_logits=self.cuda_graph_attn_logits,
attn_lse=self.cuda_graph_attn_lse,
max_extend_len=None,
num_kv_splits=self.cuda_graph_num_kv_splits,
kv_indptr=spec_info.kv_indptr,
kv_indices=spec_info.kv_indices,
qo_indptr=None,
custom_mask=None,
mask_indptr=None,
)
return
custom_mask = self.cuda_graph_custom_mask
seq_mask_len = self.num_draft_tokens * (seq_lens + self.num_draft_tokens)
mask_indptr = self.mask_indptr[: bs + 1]
mask_indptr[1 : bs + 1] = torch.cumsum(seq_mask_len, dim=0)
max_extend_len = self.num_draft_tokens
num_kv_splits = None
attn_logits = None
attn_lse = None
else:
raise ValueError(
f"Invalid forward mode: {forward_mode=} for CUDA Graph capture."
)
self.forward_metadata = ForwardMetadata(
attn_logits,
attn_lse,
max_extend_len,
num_kv_splits,
kv_indptr,
kv_indices,
qo_indptr,
custom_mask,
mask_indptr,
self.init_forward_metadata_replay_cuda_graph(
bs=bs,
req_pool_indices=req_pool_indices,
seq_lens=seq_lens,
seq_lens_sum=None,
encoder_lens=encoder_lens,
forward_mode=forward_mode,
spec_info=spec_info,
seq_lens_cpu=None,
)
self.forward_metadata = self._build_cuda_graph_forward_metadata(
bs, forward_mode, spec_info
)
def init_forward_metadata_replay_cuda_graph(
@@ -485,9 +476,7 @@ class WaveAttnBackend(AttentionBackend):
spec_info: Optional[SpecInput],
seq_lens_cpu: Optional[torch.Tensor],
):
# NOTE: encoder_lens expected to be zeros or None
if forward_mode.is_decode_or_idle():
# Update kv_indptr, kv_indices
kv_indptr = self.kv_indptr
kv_indices = self.cuda_graph_kv_indices
num_kv_splits = self.cuda_graph_num_kv_splits
@@ -510,7 +499,6 @@ class WaveAttnBackend(AttentionBackend):
num_token = spec_info.kv_indptr.shape[0] - 1
self.get_num_kv_splits(num_kv_splits[:num_token], seq_lens[:bs])
elif forward_mode.is_target_verify():
# Update qo_indptr, kv_indptr, kv_indices, custom_mask, mask_indptr
bs = len(req_pool_indices)
qo_indptr = self.qo_indptr[: bs + 1]
qo_indptr[: bs + 1] = torch.arange(