From 83a18e687d081ca800699e11765e13bf63d87182 Mon Sep 17 00:00:00 2001 From: Cheng Wan <54331508+ch-wan@users.noreply.github.com> Date: Sat, 23 May 2026 02:32:08 -0700 Subject: [PATCH] Revert "[refactor] unify cuda-graph capture/replay across attention backends (#26134)" (#26166) Co-authored-by: Cheng Wan --- .../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 | 616 ++++++++++-------- .../srt/layers/attention/wave_backend.py | 132 ++-- 5 files changed, 618 insertions(+), 461 deletions(-) diff --git a/python/sglang/srt/layers/attention/cutlass_mla_backend.py b/python/sglang/srt/layers/attention/cutlass_mla_backend.py index 22781284d..05641ea1f 100644 --- a/python/sglang/srt/layers/attention/cutlass_mla_backend.py +++ b/python/sglang/srt/layers/attention/cutlass_mla_backend.py @@ -153,22 +153,24 @@ class CutlassMLABackend(FlashInferMLAAttnBackend): forward_mode: ForwardMode, spec_info: Optional[SpecInput], ): - 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], - ) + 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], + ) else: super().init_forward_metadata_capture_cuda_graph( bs, @@ -191,11 +193,15 @@ 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[:bs], + seq_lens, 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 777a5f490..13930a752 100644 --- a/python/sglang/srt/layers/attention/flashinfer_backend.py +++ b/python/sglang/srt/layers/attention/flashinfer_backend.py @@ -559,81 +559,6 @@ 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, @@ -644,24 +569,148 @@ 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(): - for w in self.decode_cuda_graph_metadata[bs]: - w.begin_forward = partial(fast_decode_plan, w) + 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=}") def init_forward_metadata_replay_cuda_graph( self, @@ -686,7 +735,19 @@ class FlashInferAttnBackend(AttentionBackend): fixed_split_size=None, disable_split_kv=self.disable_cuda_graph_kv_split, ) - elif forward_mode.is_target_verify() or forward_mode.is_draft_extend(): + 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(): 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 a24ac11e4..61b6c49a5 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() or forward_mode.is_draft_extend(): - prefill_wrapper = BatchMLAPagedAttentionWrapper( + elif forward_mode.is_target_verify(): + verify_wrapper = BatchMLAPagedAttentionWrapper( self.workspace_buffer, use_cuda_graph=True, qo_indptr=self.cuda_graph_qo_indptr[: bs + 1], @@ -423,12 +423,34 @@ class FlashInferMLAAttnBackend(AttentionBackend): seq_lens, seq_lens_sum, prefix_lens=None, - prefill_wrapper_paged=prefill_wrapper, + prefill_wrapper_paged=verify_wrapper, use_ragged=False, spec_info=spec_info, ) - self.prefill_cuda_graph_metadata[bs] = prefill_wrapper - self.forward_metadata = PrefillMetadata(prefill_wrapper, False) + 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) else: raise ValueError(f"Invalid mode: {forward_mode=}") @@ -466,7 +488,17 @@ class FlashInferMLAAttnBackend(AttentionBackend): spec_info=spec_info, **self.fast_decode_kwargs, ) - elif forward_mode.is_target_verify() or forward_mode.is_draft_extend(): + 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(): 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 286a1399b..1ec567634 100644 --- a/python/sglang/srt/layers/attention/triton_backend.py +++ b/python/sglang/srt/layers/attention/triton_backend.py @@ -283,152 +283,11 @@ 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=self.token_to_kv_pool, - 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=self.token_to_kv_pool, - 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 @@ -438,14 +297,19 @@ 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 ) - kv_indptr = self._fill_kv_indptr_and_indices( - bs, - forward_batch.seq_lens, + create_flashinfer_kv_indices_triton[(bs,)]( + self.req_to_token, forward_batch.req_pool_indices, + forward_batch.seq_lens, + kv_indptr, + None, kv_indices, + self.req_to_token.stride(0), ) # Sliding window if ( @@ -507,14 +371,19 @@ 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( - forward_batch.seq_lens_sum, dtype=torch.int64, device=self.device + kv_indptr[-1], dtype=torch.int64, device=self.device ) - kv_indptr = self._fill_kv_indptr_and_indices( - bs, - forward_batch.seq_lens, + create_flashinfer_kv_indices_triton[(bs,)]( + self.req_to_token, forward_batch.req_pool_indices, + forward_batch.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: @@ -566,16 +435,23 @@ 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, ) - kv_indptr = self._fill_kv_indptr_and_indices( - bs, - forward_batch.extend_prefix_lens, + create_flashinfer_kv_indices_triton[(bs,)]( + self.req_to_token, forward_batch.req_pool_indices, + forward_batch.extend_prefix_lens, + kv_indptr, + None, kv_indices, + self.req_to_token.stride(0), ) # Sliding window if self.sliding_window_size is not None and self.sliding_window_size > 0: @@ -702,83 +578,6 @@ 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, @@ -790,42 +589,158 @@ 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 - # 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, + 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, + ) + ) + else: + kv_indptr, kv_indices = spec_info.kv_indptr, spec_info.kv_indices + + 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), ) - return - # 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 + 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, + ) + ) + + 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, ) def init_forward_metadata_replay_cuda_graph( @@ -841,22 +756,120 @@ class TritonAttnBackend(AttentionBackend): ): # NOTE: encoder_lens expected to be zeros or None if forward_mode.is_decode_or_idle(): - 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] + # 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), ) + 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, + ) + 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) - self._update_target_verify_buffers( - bs, seq_lens, req_pool_indices, spec_info + 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_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, + ) + ) + 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): - self._update_draft_extend_buffers(bs, seq_lens, req_pool_indices) + 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), + ) else: raise ValueError( f"Invalid forward mode: {forward_mode=} for CUDA Graph replay." @@ -1444,26 +1457,18 @@ def update_sliding_window_buffer( seq_lens, req_pool_indices, bs, - device=None, + device, token_to_kv_pool=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] - if window_kv_indices is None: - window_kv_indices = torch.empty( - window_kv_indptr[-1], dtype=torch.int64, device=device - ) + 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, @@ -1474,6 +1479,47 @@ def update_sliding_window_buffer( window_kv_indices, req_to_token.stride(0), ) + # full to swa index mapping + if hasattr(token_to_kv_pool, "translate_loc_from_full_to_swa"): + kv_last_index = window_kv_indptr[-1] + # Flush before+after: window_kv_indices is a different tensor than out_cache_loc. + token_to_kv_pool.invalidate_loc_cache() + window_kv_indices[:kv_last_index] = ( + token_to_kv_pool.translate_loc_from_full_to_swa( + window_kv_indices[:kv_last_index] + ) + ) + token_to_kv_pool.invalidate_loc_cache() + 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=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, "translate_loc_from_full_to_swa"): kv_last_index = window_kv_indptr[-1] # Flush before+after: window_kv_indices is a different tensor than out_cache_loc. diff --git a/python/sglang/srt/layers/attention/wave_backend.py b/python/sglang/srt/layers/attention/wave_backend.py index ff304a13d..ebdadc9bf 100644 --- a/python/sglang/srt/layers/attention/wave_backend.py +++ b/python/sglang/srt/layers/attention/wave_backend.py @@ -390,39 +390,6 @@ 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, @@ -435,34 +402,76 @@ class WaveAttnBackend(AttentionBackend): ): assert encoder_lens is None, "Not supported" - # 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 + 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 - 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 + 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), + ) + + 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, ) def init_forward_metadata_replay_cuda_graph( @@ -476,7 +485,9 @@ 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 @@ -499,6 +510,7 @@ 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(