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

Co-authored-by: Cheng Wan <cheng.wan@radixark.ai>

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:
+        if forward_mode.is_decode_or_idle():
-            self.init_forward_metadata_replay_cuda_graph(
+            if spec_info is None:
-                bs=bs,
+                max_seqlen_pad = self.cuda_graph_kv_indices.shape[1]
-                req_pool_indices=req_pool_indices,
+
-                seq_lens=seq_lens,
+                create_flashmla_kv_indices_triton[(bs,)](
-                seq_lens_sum=None,
+                    self.req_to_token,
-                encoder_lens=encoder_lens,
+                    req_pool_indices,
-                forward_mode=forward_mode,
+                    seq_lens,
-                spec_info=spec_info,
+                    None,
-                seq_lens_cpu=None,
+                    self.cuda_graph_kv_indices,
-            )
+                    self.req_to_token.stride(0),
-            max_seqlen_pad = self.cuda_graph_kv_indices.shape[1]
+                    self.cuda_graph_kv_indices.stride(0),
-            self.forward_metadata = CutlassMLADecodeMetadata(
+                    PAGED_SIZE=PAGE_SIZE,
-                self.cuda_graph_mla_workspace,
+                )
-                self.cuda_graph_kv_indices[:bs, :max_seqlen_pad],
+                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),

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]:
+            decode_wrappers = []
-                w.begin_forward = partial(fast_decode_plan, w)
+            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],

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():
+        elif forward_mode.is_target_verify():
-            prefill_wrapper = BatchMLAPagedAttentionWrapper(
+            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.prefill_cuda_graph_metadata[bs] = verify_wrapper
-            self.forward_metadata = PrefillMetadata(prefill_wrapper, False)
+            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],

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(
+                create_flashinfer_kv_indices_triton[(bs,)](
-                    bs,
+                    self.req_to_token,
-                    forward_batch.seq_lens,
                     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(
+            create_flashinfer_kv_indices_triton[(bs,)](
-                bs,
+                self.req_to_token,
-                forward_batch.seq_lens,
                 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(
+            create_flashinfer_kv_indices_triton[(bs,)](
-                bs,
+                self.req_to_token,
-                forward_batch.extend_prefix_lens,
                 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
+        if forward_mode.is_decode_or_idle():
-        # than the cuda-graph pool, so replay is not involved for this path.
+            if spec_info is None:
-        if forward_mode.is_decode_or_idle() and spec_info is not None:
+                kv_indptr = self.kv_indptr
-            self.forward_metadata = ForwardMetadata(
+                kv_indptr[1 : bs + 1] = torch.cumsum(seq_lens, dim=0)
-                attn_logits=self.cuda_graph_attn_logits,
+                kv_indptr = kv_indptr[: bs + 1]
-                attn_lse=self.cuda_graph_attn_lse,
+                kv_indices = self.cuda_graph_kv_indices
-                max_extend_len=None,
+                create_flashinfer_kv_indices_triton[(bs,)](
-                num_kv_splits=self.cuda_graph_num_kv_splits,
+                    self.req_to_token,
-                kv_indptr=spec_info.kv_indptr,
+                    req_pool_indices,
-                kv_indices=spec_info.kv_indices,
+                    seq_lens,
-                qo_indptr=None,
+                    kv_indptr,
-                custom_mask=None,
+                    None,
-                mask_indptr=None,
+                    kv_indices,
-                window_kv_indptr=self.window_kv_indptr,
+                    self.req_to_token.stride(0),
-                window_kv_indices=None,
+                )
-                window_num_kv_splits=None,
+                if (
-                window_kv_offsets=None,
+                    self.sliding_window_size is not None
-                swa_attn_logits=self.cuda_graph_swa_attn_logits,
+                    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
+            if self.sliding_window_size is not None and self.sliding_window_size > 0:
-        # a ForwardMetadata whose tensor fields point into the cuda-graph buffers.
+                window_kv_indices = self.cuda_graph_window_kv_indices
-        self.init_forward_metadata_replay_cuda_graph(
+                window_num_kv_splits = self.cuda_graph_window_num_kv_splits
-            bs=bs,
+                window_kv_offsets = self.cuda_graph_window_kv_offsets
-            req_pool_indices=req_pool_indices,
+                window_kv_indptr, window_kv_indices, _, window_kv_offsets[:bs] = (
-            seq_lens=seq_lens,
+                    update_sliding_window_buffer_cuda_graph(
-            seq_lens_sum=None,
+                        self.window_kv_indptr,
-            encoder_lens=encoder_lens,
+                        window_kv_indices,
-            forward_mode=forward_mode,
+                        self.req_to_token,
-            spec_info=spec_info,
+                        self.sliding_window_size,
-            seq_lens_cpu=None,
+                        seq_lens[:bs],
-        )
+                        req_pool_indices,
-        self.forward_metadata = self._build_cuda_graph_forward_metadata(
+                        bs,
-            bs, forward_mode, spec_info
+                        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."
+            # Update kv_indptr, kv_indices
-            _, _, window_kv_lens = self._update_decode_kv_buffers(
+            kv_indptr = self.kv_indptr
-                bs, seq_lens, req_pool_indices
+            kv_indices = self.cuda_graph_kv_indices
-            )
+            num_kv_splits = self.cuda_graph_num_kv_splits
-            self.get_num_kv_splits(self.cuda_graph_num_kv_splits[:bs], seq_lens[:bs])
+            if spec_info is None:
-            if window_kv_lens is not None:
+                kv_indptr[1 : bs + 1] = torch.cumsum(seq_lens[:bs], dim=0)
-                self.get_num_kv_splits(
+                kv_indptr = kv_indptr[: bs + 1]
-                    self.cuda_graph_window_num_kv_splits[:bs], window_kv_lens[:bs]
+                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(
+            qo_indptr = self.qo_indptr[: bs + 1]
-                bs, seq_lens, req_pool_indices, spec_info
+            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_indices = torch.empty(
+        window_kv_indptr[-1], dtype=torch.int64, device=device
-            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.

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
+        if forward_mode.is_decode_or_idle():
-        # the cuda-graph pool, so replay is not involved for this path.
+            if spec_info is None:
-        if forward_mode.is_decode_or_idle() and spec_info is not None:
+                kv_indptr = self.kv_indptr
-            self.forward_metadata = ForwardMetadata(
+                kv_indptr[1 : bs + 1] = torch.cumsum(seq_lens, dim=0)
-                attn_logits=self.cuda_graph_attn_logits,
+                kv_indptr = kv_indptr[: bs + 1]
-                attn_lse=self.cuda_graph_attn_lse,
+                kv_indices = self.cuda_graph_kv_indices
-                max_extend_len=None,
+                create_flashinfer_kv_indices_triton[(bs,)](
-                num_kv_splits=self.cuda_graph_num_kv_splits,
+                    self.req_to_token,
-                kv_indptr=spec_info.kv_indptr,
+                    req_pool_indices,
-                kv_indices=spec_info.kv_indices,
+                    seq_lens,
-                qo_indptr=None,
+                    kv_indptr,
-                custom_mask=None,
+                    None,
-                mask_indptr=None,
+                    kv_indices,
-            )
+                    self.req_to_token.stride(0),
-            return
+                )
+            else:
+                kv_indptr, kv_indices = spec_info.kv_indptr, spec_info.kv_indices
 
-        self.init_forward_metadata_replay_cuda_graph(
+            attn_logits = self.cuda_graph_attn_logits
-            bs=bs,
+            attn_lse = self.cuda_graph_attn_lse
-            req_pool_indices=req_pool_indices,
+            max_extend_len = None
-            seq_lens=seq_lens,
+            num_kv_splits = self.cuda_graph_num_kv_splits
-            seq_lens_sum=None,
+            qo_indptr = None
-            encoder_lens=encoder_lens,
+            custom_mask = None
-            forward_mode=forward_mode,
+            mask_indptr = None
-            spec_info=spec_info,
+        elif forward_mode.is_target_verify():
-            seq_lens_cpu=None,
+            qo_indptr = self.qo_indptr[: bs + 1]
-        )
+            qo_indptr[: bs + 1] = torch.arange(
-        self.forward_metadata = self._build_cuda_graph_forward_metadata(
+                0,
-            bs, forward_mode, spec_info
+                (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(