sglang 0.4.3.post2__py3-none-any.whl → 0.4.3.post4__py3-none-any.whl

sglang/srt/layers/attention/flashinfer_backend.py

@@ -7,28 +7,26 @@ FlashInfer is faster and Triton is easier to customize.
 Each backend supports two operators: extend (i.e. prefill with cached prefix) and decode.
 """
 
-import math
 import os
 from dataclasses import dataclass
 from enum import Enum, auto
 from functools import partial
-from typing import TYPE_CHECKING, List, Optional, Union
+from typing import TYPE_CHECKING, Callable, List, Optional, Union
 
 import torch
 import triton
-import triton.language as tl
 
 from sglang.global_config import global_config
-from sglang.srt.layers.attention import AttentionBackend
+from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
+from sglang.srt.layers.attention.utils import create_flashinfer_kv_indices_triton
 from sglang.srt.layers.dp_attention import get_attention_tp_size
-from sglang.srt.managers.schedule_batch import global_server_args_dict
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
+from sglang.srt.speculative.eagle_utils import EagleDraftInput, EagleVerifyInput
 from sglang.srt.utils import is_flashinfer_available
 
 if TYPE_CHECKING:
     from sglang.srt.layers.radix_attention import RadixAttention
     from sglang.srt.model_executor.model_runner import ModelRunner
-    from sglang.srt.speculative.spec_info import SpecInfo
 
 if is_flashinfer_available():
     from flashinfer import (
@@ -37,7 +35,7 @@ if is_flashinfer_available():
         BatchPrefillWithRaggedKVCacheWrapper,
     )
     from flashinfer.cascade import merge_state
-    from flashinfer.mla import BatchMLAPagedAttentionWrapper
+    from flashinfer.decode import _get_range_buf, get_seq_lens
 
 
 class WrapperDispatch(Enum):
@@ -47,16 +45,12 @@ class WrapperDispatch(Enum):
 
 @dataclass
 class DecodeMetadata:
-    decode_wrappers: List[
-        Union[BatchDecodeWithPagedKVCacheWrapper, BatchMLAPagedAttentionWrapper]
-    ]
+    decode_wrappers: List[BatchDecodeWithPagedKVCacheWrapper]
 
 
 @dataclass
 class PrefillMetadata:
-    prefill_wrappers: List[
-        Union[BatchPrefillWithPagedKVCacheWrapper, BatchMLAPagedAttentionWrapper]
-    ]
+    prefill_wrappers: List[BatchPrefillWithPagedKVCacheWrapper]
     use_ragged: bool
     extend_no_prefix: bool
 
@@ -73,11 +67,10 @@ class FlashInferAttnBackend(AttentionBackend):
         model_runner: ModelRunner,
         skip_prefill: bool = False,
         kv_indptr_buf: Optional[torch.Tensor] = None,
+        kv_last_page_len_buf: Optional[torch.Tensor] = None,
     ):
         super().__init__()
 
-        self.is_multimodal = model_runner.model_config.is_multimodal
-
         # Parse constants
         self.decode_use_tensor_cores = should_use_tensor_core(
             kv_cache_dtype=model_runner.kv_cache_dtype,
@@ -89,6 +82,7 @@ class FlashInferAttnBackend(AttentionBackend):
         )
         self.max_context_len = model_runner.model_config.context_len
         self.skip_prefill = skip_prefill
+        self.is_multimodal = model_runner.model_config.is_multimodal
 
         assert not (
             model_runner.sliding_window_size is not None
@@ -109,12 +103,6 @@ class FlashInferAttnBackend(AttentionBackend):
         if "Qwen2ForCausalLM" in model_runner.model_config.hf_config.architectures:
             global_config.flashinfer_workspace_size = 512 * 1024 * 1024
 
-        self.enable_flashinfer_mla = False
-        if "DeepseekV3ForCausalLM" in model_runner.model_config.hf_config.architectures:
-            if global_server_args_dict["enable_flashinfer_mla"]:
-                self.enable_flashinfer_mla = True
-                global_config.enable_flashinfer_mla = True
-
         # Allocate buffers
         global global_workspace_buffer
         if global_workspace_buffer is None:
@@ -132,24 +120,25 @@ class FlashInferAttnBackend(AttentionBackend):
                 )
                 for _ in range(self.num_wrappers)
             ]
-            if self.enable_flashinfer_mla:
-                self.qo_indptr = [
-                    torch.zeros(
-                        (max_bs + 1,), dtype=torch.int32, device=model_runner.device
-                    )
-                    for _ in range(self.num_wrappers)
-                ]
         else:
             assert self.num_wrappers == 1
             self.kv_indptr = [kv_indptr_buf]
 
-        self.kv_last_page_len = torch.ones(
-            (max_bs,), dtype=torch.int32, device=model_runner.device
-        )
-        self.qo_indptr = [
-            torch.zeros((max_bs + 1,), dtype=torch.int32, device=model_runner.device)
-            for _ in range(self.num_wrappers)
-        ]
+        if kv_last_page_len_buf is None:
+            self.kv_last_page_len = torch.ones(
+                (max_bs,), dtype=torch.int32, device=model_runner.device
+            )
+        else:
+            assert self.num_wrappers == 1
+            self.kv_last_page_len = kv_last_page_len_buf
+
+        if not self.skip_prefill:
+            self.qo_indptr = [
+                torch.zeros(
+                    (max_bs + 1,), dtype=torch.int32, device=model_runner.device
+                )
+                for _ in range(self.num_wrappers)
+            ]
 
         self.prefill_wrapper_ragged = BatchPrefillWithRaggedKVCacheWrapper(
             self.workspace_buffer, "NHD"
@@ -162,60 +151,39 @@ class FlashInferAttnBackend(AttentionBackend):
         self.decode_wrappers = []
         for _ in range(self.num_wrappers):
             if not skip_prefill:
-                if (
-                    self.enable_flashinfer_mla
-                    and not global_server_args_dict["disable_radix_cache"]
-                ):
-                    # use mla paged prefill
-                    self.prefill_wrappers_paged.append(
-                        BatchMLAPagedAttentionWrapper(
-                            self.workspace_buffer,
-                            backend="fa2",
-                        )
-                    )
-                    self.prefill_wrappers_verify.append(
-                        BatchMLAPagedAttentionWrapper(
-                            self.workspace_buffer,
-                            backend="fa2",
-                        )
-                    )
-                else:
-                    self.prefill_wrappers_paged.append(
-                        BatchPrefillWithPagedKVCacheWrapper(
-                            self.workspace_buffer,
-                            "NHD",
-                            backend="fa2",
-                        )
-                    )
-                    self.prefill_wrappers_verify.append(
-                        BatchPrefillWithPagedKVCacheWrapper(
-                            self.workspace_buffer, "NHD"
-                        )
+                self.prefill_wrappers_paged.append(
+                    BatchPrefillWithPagedKVCacheWrapper(
+                        self.workspace_buffer,
+                        "NHD",
+                        backend="fa2",
                     )
-            if self.enable_flashinfer_mla:
-                self.decode_wrappers.append(
-                    BatchMLAPagedAttentionWrapper(self.workspace_buffer, backend="fa2")
                 )
-            else:
-                self.decode_wrappers.append(
-                    BatchDecodeWithPagedKVCacheWrapper(
+                self.prefill_wrappers_verify.append(
+                    BatchPrefillWithPagedKVCacheWrapper(
                         self.workspace_buffer,
                         "NHD",
-                        use_tensor_cores=self.decode_use_tensor_cores,
                     )
                 )
+            self.decode_wrappers.append(
+                BatchDecodeWithPagedKVCacheWrapper(
+                    self.workspace_buffer,
+                    "NHD",
+                    use_tensor_cores=self.decode_use_tensor_cores,
+                )
+            )
 
         # Create indices updater
         if not skip_prefill:
             self.indices_updater_prefill = FlashInferIndicesUpdaterPrefill(
                 model_runner, self
-            )
+            )  # for verify
         self.indices_updater_decode = FlashInferIndicesUpdaterDecode(model_runner, self)
 
         # Other metadata
         self.forward_metadata: Union[PrefillMetadata, DecodeMetadata] = None
         self.decode_cuda_graph_metadata = {}
-        self.prefill_cuda_graph_metadata = {}
+        self.prefill_cuda_graph_metadata = {}  # For verify
+        self.draft_extend_cuda_graph_metadata = {}  # For draft extend
 
     def init_forward_metadata(self, forward_batch: ForwardBatch):
         if forward_batch.forward_mode.is_decode_or_idle():
@@ -259,10 +227,7 @@ class FlashInferAttnBackend(AttentionBackend):
         else:
             prefix_lens = forward_batch.extend_prefix_lens
 
-            if self.is_multimodal or (
-                self.enable_flashinfer_mla
-                and not global_server_args_dict["disable_radix_cache"]
-            ):
+            if self.is_multimodal:
                 use_ragged = False
                 extend_no_prefix = False
             else:
@@ -316,37 +281,24 @@ class FlashInferAttnBackend(AttentionBackend):
         seq_lens: torch.Tensor,
         encoder_lens: Optional[torch.Tensor],
         forward_mode: ForwardMode,
-        spec_info: Optional[SpecInfo],
+        spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]],
     ):
         if forward_mode.is_decode_or_idle():
             decode_wrappers = []
             for i in range(self.num_wrappers):
-                if self.enable_flashinfer_mla:
-                    decode_wrappers.append(
-                        BatchMLAPagedAttentionWrapper(
-                            self.workspace_buffer,
-                            use_cuda_graph=True,
-                            qo_indptr=self.qo_indptr[i][: num_tokens + 1],
-                            kv_indptr=self.kv_indptr[i][: num_tokens + 1],
-                            kv_indices=self.cuda_graph_kv_indices[i],
-                            kv_len_arr=self.kv_last_page_len[:num_tokens],
-                            backend="fa2",
-                        )
-                    )
-                else:
-                    decode_wrappers.append(
-                        BatchDecodeWithPagedKVCacheWrapper(
-                            self.workspace_buffer,
-                            "NHD",
-                            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
-                            ],
-                        )
+                decode_wrappers.append(
+                    BatchDecodeWithPagedKVCacheWrapper(
+                        self.workspace_buffer,
+                        "NHD",
+                        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,
@@ -358,6 +310,10 @@ class FlashInferAttnBackend(AttentionBackend):
             )
             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():
             prefill_wrappers = []
             for i in range(self.num_wrappers):
@@ -398,7 +354,8 @@ class FlashInferAttnBackend(AttentionBackend):
         seq_lens_sum: int,
         encoder_lens: Optional[torch.Tensor],
         forward_mode: ForwardMode,
-        spec_info: Optional[SpecInfo],
+        spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]],
+        seq_lens_cpu: Optional[torch.Tensor],
     ):
         if forward_mode.is_decode_or_idle():
             self.indices_updater_decode.update(
@@ -435,114 +392,64 @@ class FlashInferAttnBackend(AttentionBackend):
         forward_batch: ForwardBatch,
         save_kv_cache=True,
     ):
-        if global_config.enable_flashinfer_mla:
-            cache_loc = (
-                forward_batch.out_cache_loc
-                if not layer.is_cross_attention
-                else forward_batch.encoder_out_cache_loc
-            )
-
-            logits_soft_cap = layer.logit_cap
+        prefill_wrapper_paged = self.forward_metadata.prefill_wrappers[
+            self._get_wrapper_idx(layer)
+        ]
+        cache_loc = (
+            forward_batch.out_cache_loc
+            if not layer.is_cross_attention
+            else forward_batch.encoder_out_cache_loc
+        )
 
-            if global_server_args_dict["disable_radix_cache"]:
-                # use mla ragged prefill
-                o, _ = self.prefill_wrapper_ragged.forward_return_lse(
-                    q.view(-1, layer.tp_q_head_num, layer.head_dim),
-                    k.view(-1, layer.tp_k_head_num, layer.head_dim),
-                    v.view(-1, layer.tp_v_head_num, layer.v_head_dim),
-                    causal=True,
-                    sm_scale=layer.scaling,
-                    logits_soft_cap=logits_soft_cap,
-                )
+        logits_soft_cap = layer.logit_cap
 
+        if not self.forward_metadata.use_ragged:
+            if k is not None:
+                assert v is not None
                 if save_kv_cache:
                     forward_batch.token_to_kv_pool.set_kv_buffer(
-                        layer,
-                        cache_loc,
-                        k,
-                        v,
+                        layer, cache_loc, k, v, layer.k_scale, layer.v_scale
                     )
-            else:
-                # use mla paged prefill
-                prefill_wrapper_paged = self.forward_metadata.prefill_wrappers[
-                    self._get_wrapper_idx(layer)
-                ]
-                if k is not None:
-                    assert v is not None
-                    if save_kv_cache:
-                        forward_batch.token_to_kv_pool.set_kv_buffer(
-                            layer, cache_loc, k, v
-                        )
-                qall = q.view(-1, layer.tp_q_head_num, layer.head_dim)
-                k_buf = forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id)
-
-                o = prefill_wrapper_paged.run(
-                    qall[:, :, : layer.v_head_dim],
-                    qall[:, :, layer.v_head_dim :],
-                    k_buf[:, :, : layer.v_head_dim],
-                    k_buf[:, :, layer.v_head_dim :],
-                )
 
-            return o.view(-1, layer.tp_q_head_num * layer.v_head_dim)
+            o = prefill_wrapper_paged.forward(
+                q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim),
+                forward_batch.token_to_kv_pool.get_kv_buffer(layer.layer_id),
+                causal=not layer.is_cross_attention,
+                sm_scale=layer.scaling,
+                window_left=layer.sliding_window_size,
+                logits_soft_cap=logits_soft_cap,
+                k_scale=layer.k_scale,
+                v_scale=layer.v_scale,
+            )
         else:
-            prefill_wrapper_paged = self.forward_metadata.prefill_wrappers[
-                self._get_wrapper_idx(layer)
-            ]
-            cache_loc = (
-                forward_batch.out_cache_loc
-                if not layer.is_cross_attention
-                else forward_batch.encoder_out_cache_loc
+            o1, s1 = self.prefill_wrapper_ragged.forward_return_lse(
+                q.view(-1, layer.tp_q_head_num, layer.head_dim),
+                k.view(-1, layer.tp_k_head_num, layer.head_dim),
+                v.view(-1, layer.tp_v_head_num, layer.head_dim),
+                causal=True,
+                sm_scale=layer.scaling,
+                logits_soft_cap=logits_soft_cap,
             )
 
-            logits_soft_cap = layer.logit_cap
-
-            if not self.forward_metadata.use_ragged:
-                if k is not None:
-                    assert v is not None
-                    if save_kv_cache:
-                        forward_batch.token_to_kv_pool.set_kv_buffer(
-                            layer, cache_loc, k, v, layer.k_scale, layer.v_scale
-                        )
-
-                o = prefill_wrapper_paged.forward(
+            if self.forward_metadata.extend_no_prefix:
+                o = o1
+            else:
+                o2, s2 = prefill_wrapper_paged.forward_return_lse(
                     q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim),
                     forward_batch.token_to_kv_pool.get_kv_buffer(layer.layer_id),
-                    causal=not layer.is_cross_attention,
-                    sm_scale=layer.scaling,
-                    window_left=layer.sliding_window_size,
-                    logits_soft_cap=logits_soft_cap,
-                    k_scale=layer.k_scale,
-                    v_scale=layer.v_scale,
-                )
-            else:
-                o1, s1 = self.prefill_wrapper_ragged.forward_return_lse(
-                    q.view(-1, layer.tp_q_head_num, layer.head_dim),
-                    k.view(-1, layer.tp_k_head_num, layer.head_dim),
-                    v.view(-1, layer.tp_v_head_num, layer.head_dim),
-                    causal=True,
+                    causal=False,
                     sm_scale=layer.scaling,
                     logits_soft_cap=logits_soft_cap,
                 )
 
-                if self.forward_metadata.extend_no_prefix:
-                    o = o1
-                else:
-                    o2, s2 = prefill_wrapper_paged.forward_return_lse(
-                        q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim),
-                        forward_batch.token_to_kv_pool.get_kv_buffer(layer.layer_id),
-                        causal=False,
-                        sm_scale=layer.scaling,
-                        logits_soft_cap=layer.logit_cap,
-                    )
-
-                    o, _ = merge_state(o1, s1, o2, s2)
+                o, _ = merge_state(o1, s1, o2, s2)
 
-                if save_kv_cache:
-                    forward_batch.token_to_kv_pool.set_kv_buffer(
-                        layer, cache_loc, k, v, layer.k_scale, layer.v_scale
-                    )
+            if save_kv_cache:
+                forward_batch.token_to_kv_pool.set_kv_buffer(
+                    layer, cache_loc, k, v, layer.k_scale, layer.v_scale
+                )
 
-            return o.view(-1, layer.tp_q_head_num * layer.head_dim)
+        return o.view(-1, layer.tp_q_head_num * layer.head_dim)
 
     def forward_decode(
         self,
@@ -562,45 +469,23 @@ class FlashInferAttnBackend(AttentionBackend):
             else forward_batch.encoder_out_cache_loc
         )
 
-        if self.enable_flashinfer_mla:
-            if k is not None:
-                assert v is not None
-                if save_kv_cache:
-                    forward_batch.token_to_kv_pool.set_kv_buffer(
-                        layer,
-                        cache_loc,
-                        k,
-                        v,
-                    )
-            reshaped_q = q.view(-1, layer.tp_q_head_num, layer.head_dim)
-            k_buffer = forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id)
-            reshaped_k = k_buffer.view(-1, 1, layer.head_dim)
-            o = decode_wrapper.run(
-                reshaped_q[:, :, : layer.v_head_dim],
-                reshaped_q[:, :, layer.v_head_dim :],
-                reshaped_k[:, :, : layer.v_head_dim],
-                reshaped_k[:, :, layer.v_head_dim :],
-            )
-
-            return o.view(-1, layer.tp_q_head_num * layer.v_head_dim)
-        else:
-            if k is not None:
-                assert v is not None
-                if save_kv_cache:
-                    forward_batch.token_to_kv_pool.set_kv_buffer(
-                        layer, cache_loc, k, v, layer.k_scale, layer.v_scale
-                    )
+        if k is not None:
+            assert v is not None
+            if save_kv_cache:
+                forward_batch.token_to_kv_pool.set_kv_buffer(
+                    layer, cache_loc, k, v, layer.k_scale, layer.v_scale
+                )
 
-            o = decode_wrapper.forward(
-                q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim),
-                forward_batch.token_to_kv_pool.get_kv_buffer(layer.layer_id),
-                sm_scale=layer.scaling,
-                logits_soft_cap=layer.logit_cap,
-                k_scale=layer.k_scale,
-                v_scale=layer.v_scale,
-            )
+        o = decode_wrapper.forward(
+            q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim),
+            forward_batch.token_to_kv_pool.get_kv_buffer(layer.layer_id),
+            sm_scale=layer.scaling,
+            logits_soft_cap=layer.logit_cap,
+            k_scale=layer.k_scale,
+            v_scale=layer.v_scale,
+        )
 
-            return o.view(-1, layer.tp_q_head_num * layer.head_dim)
+        return o.view(-1, layer.tp_q_head_num * layer.head_dim)
 
     def _get_wrapper_idx(self, layer: RadixAttention):
         if self.num_wrappers == 1:
@@ -648,11 +533,9 @@ class FlashInferIndicesUpdaterDecode:
         req_pool_indices: torch.Tensor,
         seq_lens: torch.Tensor,
         seq_lens_sum: int,
-        decode_wrappers: List[
-            Union[BatchDecodeWithPagedKVCacheWrapper, BatchMLAPagedAttentionWrapper]
-        ],
+        decode_wrappers: List[BatchDecodeWithPagedKVCacheWrapper],
         encoder_lens: Optional[torch.Tensor],
-        spec_info: Optional[SpecInfo],
+        spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]],
     ):
         # Keep the signature for type checking. It will be assigned during runtime.
         raise NotImplementedError()
@@ -662,11 +545,9 @@ class FlashInferIndicesUpdaterDecode:
         req_pool_indices: torch.Tensor,
         seq_lens: torch.Tensor,
         seq_lens_sum: int,
-        decode_wrappers: List[
-            Union[BatchDecodeWithPagedKVCacheWrapper, BatchMLAPagedAttentionWrapper]
-        ],
+        decode_wrappers: List[BatchDecodeWithPagedKVCacheWrapper],
         encoder_lens: Optional[torch.Tensor],
-        spec_info: Optional[SpecInfo],
+        spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]],
     ):
         decode_wrappers = decode_wrappers or self.decode_wrappers
         self.call_begin_forward(
@@ -686,7 +567,7 @@ class FlashInferIndicesUpdaterDecode:
         seq_lens_sum: int,
         decode_wrappers: List[BatchDecodeWithPagedKVCacheWrapper],
         encoder_lens: Optional[torch.Tensor],
-        spec_info: Optional[SpecInfo],
+        spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]],
     ):
         for wrapper_id in range(2):
             if wrapper_id == 0:
@@ -720,7 +601,7 @@ class FlashInferIndicesUpdaterDecode:
         seq_lens_sum: int,
         decode_wrappers: List[BatchDecodeWithPagedKVCacheWrapper],
         encoder_lens: Optional[torch.Tensor],
-        spec_info: Optional[SpecInfo],
+        spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]],
     ):
         for wrapper_id in range(2):
             if wrapper_id == 0:
@@ -745,23 +626,27 @@ class FlashInferIndicesUpdaterDecode:
 
     def call_begin_forward(
         self,
-        wrapper: Union[
-            BatchDecodeWithPagedKVCacheWrapper, BatchMLAPagedAttentionWrapper
-        ],
+        wrapper: BatchDecodeWithPagedKVCacheWrapper,
         req_pool_indices: torch.Tensor,
         paged_kernel_lens: torch.Tensor,
         paged_kernel_lens_sum: int,
         kv_indptr: torch.Tensor,
         kv_start_idx: torch.Tensor,
-        spec_info: Optional[SpecInfo],
+        spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]],
     ):
         if spec_info is None:
             bs = len(req_pool_indices)
             kv_indptr[1 : bs + 1] = torch.cumsum(paged_kernel_lens, dim=0)
             kv_indptr = kv_indptr[: bs + 1]
-            kv_indices = torch.empty(
-                paged_kernel_lens_sum, dtype=torch.int32, device="cuda"
-            )
+
+            if wrapper.is_cuda_graph_enabled:
+                # Directly write to the cuda graph input buffer
+                kv_indices = wrapper._paged_kv_indices_buf
+            else:
+                kv_indices = torch.empty(
+                    paged_kernel_lens_sum, dtype=torch.int32, device="cuda"
+                )
+
             create_flashinfer_kv_indices_triton[(bs,)](
                 self.req_to_token,
                 req_pool_indices,
@@ -775,37 +660,18 @@ class FlashInferIndicesUpdaterDecode:
             kv_indptr, kv_indices = spec_info.kv_indptr, spec_info.kv_indices
             bs = kv_indptr.shape[0] - 1
 
-        if global_config.enable_flashinfer_mla:
-            sm_scale = 1.0 / math.sqrt(192)
-            q_indptr = torch.arange(0, bs + 1).to(0).int()
-            kv_lens = paged_kernel_lens.to(torch.int32)
-            wrapper.plan(
-                q_indptr,
-                kv_indptr,
-                kv_indices,
-                kv_lens,
-                self.num_qo_heads,
-                512,
-                64,
-                1,
-                False,
-                sm_scale,
-                self.data_type,
-                self.data_type,
-            )
-        else:
-            wrapper.begin_forward(
-                kv_indptr,
-                kv_indices,
-                self.kv_last_page_len[:bs],
-                self.num_qo_heads,
-                self.num_kv_heads,
-                self.head_dim,
-                1,
-                data_type=self.data_type,
-                q_data_type=self.q_data_type,
-                non_blocking=True,
-            )
+        wrapper.begin_forward(
+            kv_indptr,
+            kv_indices,
+            self.kv_last_page_len[:bs],
+            self.num_qo_heads,
+            self.num_kv_heads,
+            self.head_dim,
+            1,
+            data_type=self.data_type,
+            q_data_type=self.q_data_type,
+            non_blocking=True,
+        )
 
 
 class FlashInferIndicesUpdaterPrefill:
@@ -841,32 +707,28 @@ class FlashInferIndicesUpdaterPrefill:
 
     def update(
         self,
-        req_pool_indices: torch.Tnesor,
+        req_pool_indices: torch.Tensor,
         seq_lens: torch.Tensor,
         seq_lens_sum: int,
         prefix_lens: torch.Tensor,
-        prefill_wrappers: List[
-            Union[BatchPrefillWithPagedKVCacheWrapper, BatchMLAPagedAttentionWrapper]
-        ],
+        prefill_wrappers: List[BatchPrefillWithPagedKVCacheWrapper],
         use_ragged: bool,
         encoder_lens: Optional[torch.Tensor],
-        spec_info: Optional[SpecInfo],
+        spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]],
     ):
         # Keep the signature for type checking. It will be assigned during runtime.
         raise NotImplementedError()
 
     def update_single_wrapper(
         self,
-        req_pool_indices: torch.Tnesor,
+        req_pool_indices: torch.Tensor,
         seq_lens: torch.Tensor,
         seq_lens_sum: int,
         prefix_lens: torch.Tensor,
-        prefill_wrappers: List[
-            Union[BatchPrefillWithPagedKVCacheWrapper, BatchMLAPagedAttentionWrapper]
-        ],
+        prefill_wrappers: List[BatchPrefillWithPagedKVCacheWrapper],
         use_ragged: bool,
         encoder_lens: Optional[torch.Tensor],
-        spec_info: Optional[SpecInfo],
+        spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]],
     ):
         if use_ragged:
             paged_kernel_lens = prefix_lens
@@ -899,7 +761,7 @@ class FlashInferIndicesUpdaterPrefill:
         prefill_wrappers: List[BatchPrefillWithPagedKVCacheWrapper],
         use_ragged: bool,
         encoder_lens: Optional[torch.Tensor],
-        spec_info: Optional[SpecInfo],
+        spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]],
     ):
         for wrapper_id in range(2):
             if wrapper_id == 0:
@@ -940,7 +802,7 @@ class FlashInferIndicesUpdaterPrefill:
         prefill_wrappers: List[BatchPrefillWithPagedKVCacheWrapper],
         use_ragged: bool,
         encoder_lens: Optional[torch.Tensor],
-        spec_info: Optional[SpecInfo],
+        spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]],
     ):
         for wrapper_id in range(2):
             if wrapper_id == 0:
@@ -972,9 +834,7 @@ class FlashInferIndicesUpdaterPrefill:
     def call_begin_forward(
         self,
         wrapper_ragged: BatchPrefillWithRaggedKVCacheWrapper,
-        wrapper_paged: Union[
-            BatchPrefillWithPagedKVCacheWrapper, BatchMLAPagedAttentionWrapper
-        ],
+        wrapper_paged: BatchPrefillWithPagedKVCacheWrapper,
         req_pool_indices: torch.Tensor,
         paged_kernel_lens: torch.Tensor,
         paged_kernel_lens_sum: int,
@@ -984,10 +844,11 @@ class FlashInferIndicesUpdaterPrefill:
         kv_indptr: torch.Tensor,
         qo_indptr: torch.Tensor,
         use_ragged: bool,
-        spec_info: Optional[SpecInfo],
+        spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]],
     ):
-        bs = len(req_pool_indices)
+        bs = len(seq_lens)
         if spec_info is None:
+            assert len(seq_lens) == len(req_pool_indices)
             # Normal extend
             kv_indptr[1 : bs + 1] = torch.cumsum(paged_kernel_lens, dim=0)
             kv_indptr = kv_indptr[: bs + 1]
@@ -1005,77 +866,54 @@ class FlashInferIndicesUpdaterPrefill:
                 kv_indices,
                 self.req_to_token.shape[1],
             )
-
             qo_indptr[1 : bs + 1] = torch.cumsum(seq_lens - prefix_lens, dim=0)
             qo_indptr = qo_indptr[: bs + 1]
             custom_mask = None
         else:
+            assert isinstance(spec_info, EagleDraftInput) or isinstance(
+                spec_info, EagleVerifyInput
+            )
             kv_indices, kv_indptr, qo_indptr, custom_mask = (
                 spec_info.generate_attn_arg_prefill(
                     req_pool_indices,
                     paged_kernel_lens,
+                    paged_kernel_lens_sum,
                     self.req_to_token,
                 )
             )
 
         # extend part
         if use_ragged:
-            if global_config.enable_flashinfer_mla:
-                wrapper_ragged.begin_forward(
-                    qo_indptr=qo_indptr,
-                    kv_indptr=qo_indptr,
-                    num_qo_heads=self.num_qo_heads,
-                    num_kv_heads=self.num_kv_heads,
-                    head_dim_qk=192,
-                    head_dim_vo=128,
-                    q_data_type=self.q_data_type,
-                )
-            else:
-                wrapper_ragged.begin_forward(
-                    qo_indptr,
-                    qo_indptr,
-                    self.num_qo_heads,
-                    self.num_kv_heads,
-                    self.head_dim,
-                    q_data_type=self.q_data_type,
-                )
-
-        if not global_config.enable_flashinfer_mla:
-            # cached part
-            wrapper_paged.begin_forward(
+            wrapper_ragged.begin_forward(
+                qo_indptr,
                 qo_indptr,
-                kv_indptr,
-                kv_indices,
-                self.kv_last_page_len[:bs],
                 self.num_qo_heads,
                 self.num_kv_heads,
                 self.head_dim,
-                1,
                 q_data_type=self.q_data_type,
-                custom_mask=custom_mask,
-                non_blocking=True,
-            )
-        elif (
-            global_config.enable_flashinfer_mla
-            and not global_server_args_dict["disable_radix_cache"]
-        ):
-            # mla paged prefill
-            kv_len_arr = kv_indptr[1:] - kv_indptr[:-1]
-            wrapper_paged.plan(
-                qo_indptr,
-                kv_indptr,
-                kv_indices,
-                kv_len_arr,
-                self.num_qo_heads,
-                512,
-                64,
-                1,
-                True,
-                1 / math.sqrt(192),
-                self.data_type,
-                self.data_type,
             )
 
+        # cached part
+        wrapper_paged.begin_forward(
+            qo_indptr,
+            kv_indptr,
+            kv_indices,
+            self.kv_last_page_len[:bs],
+            self.num_qo_heads,
+            self.num_kv_heads,
+            self.head_dim,
+            1,
+            q_data_type=self.q_data_type,
+            kv_data_type=self.data_type,
+            custom_mask=custom_mask,
+            non_blocking=True,
+        )
+
+
+# Use as a fast path to override the indptr in flashinfer's plan function
+# This is used to remove some host-to-device copy overhead.
+global global_override_indptr_cpu
+
 
 class FlashInferMultiStepDraftBackend:
     """
@@ -1094,7 +932,8 @@ class FlashInferMultiStepDraftBackend:
         self.topk = topk
         self.speculative_num_steps = speculative_num_steps
         self.generate_draft_decode_kv_indices = generate_draft_decode_kv_indices
-        max_bs = model_runner.req_to_token_pool.size
+
+        max_bs = model_runner.req_to_token_pool.size * self.topk
         self.kv_indptr = torch.zeros(
             (
                 self.speculative_num_steps,
@@ -1103,6 +942,9 @@ class FlashInferMultiStepDraftBackend:
             dtype=torch.int32,
             device=model_runner.device,
         )
+        self.kv_last_page_len = torch.ones(
+            (max_bs,), dtype=torch.int32, device=model_runner.device
+        )
         self.attn_backends = []
         for i in range(self.speculative_num_steps):
             self.attn_backends.append(
@@ -1110,14 +952,20 @@ class FlashInferMultiStepDraftBackend:
                     model_runner,
                     skip_prefill=True,
                     kv_indptr_buf=self.kv_indptr[i],
+                    kv_last_page_len_buf=self.kv_last_page_len,
                 )
             )
+
         self.max_context_len = self.attn_backends[0].max_context_len
+
         # Cached variables for generate_draft_decode_kv_indices
         self.pool_len = model_runner.req_to_token_pool.req_to_token.shape[1]
 
     def common_template(
-        self, forward_batch: ForwardBatch, kv_indices_buffer: torch.Tensor, call_fn: int
+        self,
+        forward_batch: ForwardBatch,
+        kv_indices_buffer: torch.Tensor,
+        call_fn: Callable,
     ):
         num_seqs = forward_batch.batch_size
         bs = self.topk * num_seqs
@@ -1142,13 +990,23 @@ class FlashInferMultiStepDraftBackend:
             triton.next_power_of_2(bs),
         )
 
+        assert forward_batch.spec_info is not None
+        assert isinstance(forward_batch.spec_info, EagleDraftInput)
+
+        # Copy the kv_indptr once to avoid multiple device-to-host copies in flashinfer's plan.
+        indptr_cpu_whole = self.kv_indptr[:, : bs + 1].cpu()
+        global global_override_indptr_cpu
+
         for i in range(self.speculative_num_steps - 1):
             forward_batch.spec_info.kv_indptr = self.kv_indptr[i, : bs + 1]
             forward_batch.spec_info.kv_indices = kv_indices_buffer[i][
                 : seq_lens_sum * self.topk + bs * (i + 1)
             ]
+            global_override_indptr_cpu = indptr_cpu_whole[i]
             call_fn(i, forward_batch)
 
+        global_override_indptr_cpu = None
+
     def init_forward_metadata(self, forward_batch: ForwardBatch):
         kv_indices = torch.zeros(
             (
@@ -1160,6 +1018,8 @@ class FlashInferMultiStepDraftBackend:
         )
 
         def call_fn(i, forward_batch):
+            assert forward_batch.spec_info is not None
+            assert isinstance(forward_batch.spec_info, EagleDraftInput)
             forward_batch.spec_info.kv_indptr = (
                 forward_batch.spec_info.kv_indptr.clone()
             )
@@ -1176,6 +1036,7 @@ class FlashInferMultiStepDraftBackend:
             dtype=torch.int32,
             device="cuda",
         )
+
         for i in range(self.speculative_num_steps):
             self.attn_backends[i].init_cuda_graph_state(
                 max_bs, kv_indices_buf=self.cuda_graph_kv_indices[i]
@@ -1192,65 +1053,27 @@ class FlashInferMultiStepDraftBackend:
                 forward_mode=ForwardMode.DECODE,
                 spec_info=forward_batch.spec_info,
             )
-            decode_wrapper = self.attn_backends[i].decode_cuda_graph_metadata[
-                forward_batch.batch_size
-            ][0]
-            decode_wrapper.begin_forward = partial(fast_decode_plan, decode_wrapper)
 
         self.common_template(forward_batch, self.cuda_graph_kv_indices, call_fn)
 
-    def init_forward_metadata_replay_cuda_graph(self, forward_batch):
+    def init_forward_metadata_replay_cuda_graph(
+        self, forward_batch: ForwardBatch, bs: int
+    ):
         def call_fn(i, forward_batch):
             self.attn_backends[i].init_forward_metadata_replay_cuda_graph(
-                forward_batch.batch_size,
+                bs,
                 forward_batch.req_pool_indices,
                 forward_batch.seq_lens,
                 seq_lens_sum=-1,
                 encoder_lens=None,
                 forward_mode=ForwardMode.DECODE,
                 spec_info=forward_batch.spec_info,
+                seq_lens_cpu=None,
             )
 
         self.common_template(forward_batch, self.cuda_graph_kv_indices, call_fn)
 
 
-@triton.jit
-def create_flashinfer_kv_indices_triton(
-    req_to_token_ptr,  # [max_batch, max_context_len]
-    req_pool_indices_ptr,
-    page_kernel_lens_ptr,
-    kv_indptr,
-    kv_start_idx,
-    kv_indices_ptr,
-    req_to_token_ptr_stride: tl.constexpr,
-):
-    BLOCK_SIZE: tl.constexpr = 512
-    pid = tl.program_id(axis=0)
-
-    req_pool_index = tl.load(req_pool_indices_ptr + pid)
-    kv_indices_offset = tl.load(kv_indptr + pid)
-
-    kv_start = 0
-    kv_end = 0
-    if kv_start_idx:
-        kv_start = tl.load(kv_start_idx + pid).to(tl.int32)
-        kv_end = kv_start
-    kv_end += tl.load(page_kernel_lens_ptr + pid).to(tl.int32)
-
-    num_loop = tl.cdiv(kv_end - kv_start, BLOCK_SIZE)
-    for i in range(num_loop):
-        offset = tl.arange(0, BLOCK_SIZE) + i * BLOCK_SIZE
-        mask = offset < kv_end - kv_start
-        data = tl.load(
-            req_to_token_ptr
-            + req_pool_index * req_to_token_ptr_stride
-            + kv_start
-            + offset,
-            mask=mask,
-        )
-        tl.store(kv_indices_ptr + kv_indices_offset + offset, data, mask=mask)
-
-
 def should_use_tensor_core(
     kv_cache_dtype: torch.dtype,
     num_attention_heads: int,
@@ -1272,6 +1095,21 @@ def should_use_tensor_core(
     if env_override is not None:
         return env_override.lower() == "true"
 
+    # Try to use _grouped_size_compiled_for_decode_kernels if available
+    # This is for flashinfer <=0.1.6. Otherwise, there is an accuracy bug
+    try:
+        from flashinfer.decode import _grouped_size_compiled_for_decode_kernels
+
+        if not _grouped_size_compiled_for_decode_kernels(
+            num_attention_heads,
+            num_kv_heads,
+        ):
+            return True
+        else:
+            return False
+    except (ImportError, AttributeError):
+        pass
+
     # Calculate GQA group size
     gqa_group_size = num_attention_heads // num_kv_heads
 
@@ -1284,6 +1122,11 @@ def should_use_tensor_core(
         return False
 
 
+# Use as a fast path to override the indptr in flashinfer's plan function
+# This is used to remove some host-to-device copy overhead.
+global_override_indptr_cpu = None
+
+
 def fast_decode_plan(
     self,
     indptr: torch.Tensor,
@@ -1301,12 +1144,21 @@ def fast_decode_plan(
     sm_scale: Optional[float] = None,
     rope_scale: Optional[float] = None,
     rope_theta: Optional[float] = None,
-    **kwargs,
+    non_blocking: bool = True,
 ) -> None:
-    """A faster version of BatchDecodeWithPagedKVCacheWrapper::plan used for FlashInferMultiStepDraftBackend."""
+    """
+    A faster version of BatchDecodeWithPagedKVCacheWrapper::plan used for FlashInferMultiStepDraftBackend.
+    Modifications:
+    - Remove unnecessary device-to-device copy for the cuda graph buffers.
+    - Remove unnecessary host-to-device copy for the metadata buffers.
+    """
     batch_size = len(last_page_len)
     if logits_soft_cap is None:
         logits_soft_cap = 0.0
+
+    if self.use_tensor_cores:
+        qo_indptr_host = _get_range_buf(batch_size + 1, "cpu")
+
     if self.is_cuda_graph_enabled:
         if batch_size != self._fixed_batch_size:
             raise ValueError(
@@ -1319,13 +1171,20 @@ def fast_decode_plan(
             raise ValueError(
                 "The size of indices should be less than or equal to the allocated buffer"
             )
+        # Skip these copies because we directly write to them during prepartion
+        # self._paged_kv_indptr_buf.copy_(indptr)
+        # self._paged_kv_indices_buf[: len(indices)] = indices
+        # self._paged_kv_last_page_len_buf.copy_(last_page_len)
     else:
         self._paged_kv_indptr_buf = indptr
         self._paged_kv_indices_buf = indices
         self._paged_kv_last_page_len_buf = last_page_len
+        self._qo_indptr_buf = qo_indptr_host.to(self.device, non_blocking=non_blocking)
+
     # NOTE(Zihao): the following tensors acts as placeholder to pass dtype info
     if not q_data_type:
         q_data_type = data_type
+
     if not hasattr(self, "empty_q_data"):
         self.empty_q_data = torch.empty(
             0,
@@ -1342,27 +1201,56 @@ def fast_decode_plan(
             ),
         )
         self.last_page_len = torch.ones(32768, dtype=torch.int32)
-    empty_q_data = self.empty_q_data
-    empty_kv_cache = self.empty_kv_cache
-    stream = torch.cuda.current_stream()
-    self._cached_module.plan(
-        self._float_workspace_buffer,
-        self._int_workspace_buffer,
-        self._pin_memory_int_workspace_buffer,
-        indptr.to("cpu"),
-        batch_size,
-        num_qo_heads,
-        num_kv_heads,
-        page_size,
-        self.is_cuda_graph_enabled,
-        window_left,
-        logits_soft_cap,
-        head_dim,
-        head_dim,
-        empty_q_data,
-        empty_kv_cache,
-        stream.cuda_stream,
+
+    indptr_host = (
+        global_override_indptr_cpu
+        if global_override_indptr_cpu is not None
+        else indptr.cpu()
     )
+
+    if self.use_tensor_cores:
+        kv_lens_arr_host = get_seq_lens(
+            indptr_host, self.last_page_len[:batch_size], page_size
+        )
+
+        self._plan_info = self._cached_module.plan(
+            self._float_workspace_buffer,
+            self._int_workspace_buffer,
+            self._pin_memory_int_workspace_buffer,
+            qo_indptr_host,
+            indptr_host,
+            kv_lens_arr_host,
+            batch_size,  # total_num_rows
+            batch_size,
+            num_qo_heads,
+            num_kv_heads,
+            page_size,
+            self.is_cuda_graph_enabled,
+            head_dim,
+            head_dim,
+            False,  # causal
+            torch.cuda.current_stream().cuda_stream,
+        )
+    else:
+        self._plan_info = self._cached_module.plan(
+            self._float_workspace_buffer,
+            self._int_workspace_buffer,
+            self._pin_memory_int_workspace_buffer,
+            indptr_host,
+            batch_size,
+            num_qo_heads,
+            num_kv_heads,
+            page_size,
+            self.is_cuda_graph_enabled,
+            window_left,
+            logits_soft_cap,
+            head_dim,
+            head_dim,
+            self.empty_q_data,
+            self.empty_kv_cache,
+            torch.cuda.current_stream().cuda_stream,
+        )
+
     self._pos_encoding_mode = pos_encoding_mode
     self._window_left = window_left
     self._logits_soft_cap = logits_soft_cap