From d226f7566945eceb5713c0068ff83cfbda11b84c Mon Sep 17 00:00:00 2001 From: Cheng Wan <54331508+ch-wan@users.noreply.github.com> Date: Fri, 22 May 2026 18:51:13 -0700 Subject: [PATCH] [refactor] unify cuda-graph capture/replay across attention backends (#26134) Co-authored-by: Cheng Wan Co-authored-by: Claude Sonnet 4.6 (1M context) --- .../layers/attention/cutlass_mla_backend.py | 40 +- .../layers/attention/flashinfer_backend.py | 247 +++---- .../attention/flashinfer_mla_backend.py | 44 +- .../srt/layers/attention/triton_backend.py | 615 ++++++++---------- .../srt/layers/attention/wave_backend.py | 130 ++-- 5 files changed, 461 insertions(+), 615 deletions(-) diff --git a/python/sglang/srt/layers/attention/cutlass_mla_backend.py b/python/sglang/srt/layers/attention/cutlass_mla_backend.py index 05641ea1f..22781284d 100644 --- a/python/sglang/srt/layers/attention/cutlass_mla_backend.py +++ b/python/sglang/srt/layers/attention/cutlass_mla_backend.py @@ -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), diff --git a/python/sglang/srt/layers/attention/flashinfer_backend.py b/python/sglang/srt/layers/attention/flashinfer_backend.py index 13930a752..777a5f490 100644 --- a/python/sglang/srt/layers/attention/flashinfer_backend.py +++ b/python/sglang/srt/layers/attention/flashinfer_backend.py @@ -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], diff --git a/python/sglang/srt/layers/attention/flashinfer_mla_backend.py b/python/sglang/srt/layers/attention/flashinfer_mla_backend.py index 61b6c49a5..a24ac11e4 100644 --- a/python/sglang/srt/layers/attention/flashinfer_mla_backend.py +++ b/python/sglang/srt/layers/attention/flashinfer_mla_backend.py @@ -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], diff --git a/python/sglang/srt/layers/attention/triton_backend.py b/python/sglang/srt/layers/attention/triton_backend.py index ea727713f..4fd855572 100644 --- a/python/sglang/srt/layers/attention/triton_backend.py +++ b/python/sglang/srt/layers/attention/triton_backend.py @@ -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] = ( diff --git a/python/sglang/srt/layers/attention/wave_backend.py b/python/sglang/srt/layers/attention/wave_backend.py index ebdadc9bf..ff304a13d 100644 --- a/python/sglang/srt/layers/attention/wave_backend.py +++ b/python/sglang/srt/layers/attention/wave_backend.py @@ -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(