sglang 0.2.15__py3-none-any.whl → 0.3.1__py3-none-any.whl

sglang/srt/layers/flashinfer_utils.py

@@ -0,0 +1,235 @@
+import torch
+import triton
+import triton.language as tl
+
+
+@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,
+    max_context_len: 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)
+
+    req_to_token_ptr += req_pool_index * max_context_len
+    kv_indices_ptr += kv_indices_offset
+
+    ld_offset = kv_start + tl.arange(0, BLOCK_SIZE)
+    st_offset = tl.arange(0, BLOCK_SIZE)
+    num_loop = tl.cdiv(kv_end - kv_start, BLOCK_SIZE)
+    for _ in range(num_loop):
+        mask = ld_offset < kv_end
+        data = tl.load(req_to_token_ptr + ld_offset, mask=mask)
+        tl.store(kv_indices_ptr + st_offset, data, mask=mask)
+        ld_offset += BLOCK_SIZE
+        st_offset += BLOCK_SIZE
+
+
+class FlashinferUpdater:
+    def __init__(
+        self,
+        forward_mode,
+        model_runner,
+        req_pool_indices,
+        seq_lens,
+        prefix_lens,
+        decode_wrapper=None,
+        use_ragged=False,
+    ):
+        self.forward_mode = forward_mode
+        self.model_runner = model_runner
+        self.req_pool_indices = req_pool_indices
+        self.seq_lens = seq_lens
+        self.prefix_lens = prefix_lens
+        self.use_ragged = use_ragged
+
+        self.num_qo_heads = (
+            model_runner.model_config.num_attention_heads // model_runner.tp_size
+        )
+        self.num_kv_heads = model_runner.model_config.get_num_kv_heads(
+            model_runner.tp_size
+        )
+        self.head_dim = model_runner.model_config.head_dim
+        self.batch_size = len(req_pool_indices)
+
+        self.decode_wrapper = (
+            decode_wrapper or self.model_runner.attn_backend.decode_wrapper
+        )
+        self.prefill_wrapper_ragged = (
+            self.model_runner.attn_backend.prefill_wrapper_ragged
+        )
+        self.prefill_wrapper_paged = (
+            self.model_runner.attn_backend.prefill_wrapper_paged
+        )
+
+        self.kv_last_page_len = torch.ones(
+            (self.batch_size,), dtype=torch.int32, device="cuda"
+        )
+
+    def _init_indices_no_sliding_window(self):
+        if self.use_ragged:
+            paged_kernel_lens = self.prefix_lens
+        else:
+            paged_kernel_lens = self.seq_lens
+
+        self.kv_indptr = torch.zeros(
+            (self.batch_size + 1,), dtype=torch.int32, device="cuda"
+        )
+        self.kv_indptr[1:] = torch.cumsum(paged_kernel_lens, dim=0)
+        self.kv_indices = torch.empty(
+            self.kv_indptr[-1], dtype=torch.int32, device="cuda"
+        )
+
+        create_flashinfer_kv_indices_triton[(self.batch_size,)](
+            self.model_runner.req_to_token_pool.req_to_token,
+            self.req_pool_indices,
+            paged_kernel_lens,
+            self.kv_indptr,
+            None,
+            self.kv_indices,
+            self.model_runner.req_to_token_pool.req_to_token.size(1),
+        )
+
+    def _init_indices_sliding_window(self, wrapper_id):
+        if wrapper_id == 0:
+            # window attention use paged only
+            if self.forward_mode.is_decode():
+                paged_kernel_lens = torch.minimum(
+                    self.seq_lens,
+                    torch.tensor(self.model_runner.sliding_window_size + 1),
+                )
+            else:
+                paged_kernel_lens = torch.minimum(
+                    self.seq_lens,
+                    torch.tensor(self.model_runner.sliding_window_size)
+                    + self.seq_lens
+                    - self.prefix_lens,
+                )
+        else:
+            # full attention
+            paged_kernel_lens = self.seq_lens
+
+        kv_start_idx = self.seq_lens - paged_kernel_lens
+        self.kv_indptr = torch.zeros(
+            (self.batch_size + 1,), dtype=torch.int32, device="cuda"
+        )
+        self.kv_indptr[1:] = torch.cumsum(paged_kernel_lens, dim=0)
+        self.kv_indices = torch.empty(
+            self.kv_indptr[-1], dtype=torch.int32, device="cuda"
+        )
+        create_flashinfer_kv_indices_triton[(self.batch_size,)](
+            self.model_runner.req_to_token_pool.req_to_token,
+            self.req_pool_indices,
+            paged_kernel_lens,
+            self.kv_indptr,
+            kv_start_idx,
+            self.kv_indices,
+            self.model_runner.req_to_token_pool.req_to_token.size(1),
+        )
+
+    def _update_decode_indices(self, decode_wrapper):
+        decode_wrapper.end_forward()
+        decode_wrapper.begin_forward(
+            self.kv_indptr,
+            self.kv_indices,
+            self.kv_last_page_len,
+            self.num_qo_heads,
+            self.num_kv_heads,
+            self.head_dim,
+            1,
+            data_type=self.model_runner.kv_cache_dtype,
+            q_data_type=self.model_runner.dtype,
+        )
+
+    def _update_extend_indices(self, ragged_wrapper, paged_wrapper):
+        # extend part
+        qo_indptr = torch.zeros(
+            (self.batch_size + 1,), dtype=torch.int32, device="cuda"
+        )
+        qo_indptr[1:] = torch.cumsum(self.seq_lens - self.prefix_lens, dim=0)
+
+        if self.use_ragged:
+            ragged_wrapper.end_forward()
+            ragged_wrapper.begin_forward(
+                qo_indptr,
+                qo_indptr,
+                self.num_qo_heads,
+                self.num_kv_heads,
+                self.head_dim,
+            )
+
+        # cached part
+        paged_wrapper.end_forward()
+        paged_wrapper.begin_forward(
+            qo_indptr,
+            self.kv_indptr,
+            self.kv_indices,
+            self.kv_last_page_len,
+            self.num_qo_heads,
+            self.num_kv_heads,
+            self.head_dim,
+            1,
+        )
+
+    def update_indices_no_sliding_window(self):
+        self._init_indices_no_sliding_window()
+
+        if self.forward_mode.is_decode():
+            self._update_decode_indices(self.decode_wrapper)
+        else:
+            self._update_extend_indices(
+                self.prefill_wrapper_ragged,
+                self.prefill_wrapper_paged,
+            )
+
+    def update_indices_sliding_window(self):
+        assert self.use_ragged is False
+
+        for wrapper_id in range(2):
+            self._init_indices_sliding_window(wrapper_id)
+            if self.forward_mode.is_decode():
+                self._update_decode_indices(self.decode_wrapper[wrapper_id])
+            else:
+                self._update_extend_indices(
+                    None,
+                    self.prefill_wrapper_paged[wrapper_id],
+                )
+
+
+def update_flashinfer_indices(
+    forward_mode,
+    model_runner,
+    req_pool_indices,
+    seq_lens,
+    prefix_lens,
+    decode_wrapper=None,
+    use_ragged=False,
+):
+    updater = FlashinferUpdater(
+        forward_mode,
+        model_runner,
+        req_pool_indices,
+        seq_lens,
+        prefix_lens,
+        decode_wrapper,
+        use_ragged,
+    )
+
+    if model_runner.sliding_window_size is None:
+        updater.update_indices_no_sliding_window()
+    else:
+        updater.update_indices_sliding_window()

sglang/srt/layers/logits_processor.py

@@ -37,7 +37,7 @@ class LogitsProcessorOutput:
 
     # The normlaized logprobs of prompts.  shape: [#seq]
     normalized_prompt_logprobs: torch.Tensor
-    # The logprobs of input tokens.      shape: [#token, vocab_size]
+    # The logprobs of input tokens.        shape: [#token, vocab_size]
     input_token_logprobs: torch.Tensor
 
     # The logprob and id of the top-k tokens in input positions.  shape [#seq, #token, k] of Tuple(logprob, token_id)
@@ -49,25 +49,39 @@ class LogitsProcessorOutput:
 @dataclasses.dataclass
 class LogitsMetadata:
     forward_mode: ForwardMode
+    top_logprobs_nums: Optional[List[int]]
+
     return_logprob: bool = False
+    return_top_logprob: bool = False
 
     extend_seq_lens: Optional[torch.Tensor] = None
-    extend_start_loc: Optional[torch.Tensor] = None
-    top_logprobs_nums: Optional[List[int]] = None
+    extend_seq_lens_cpu: Optional[List[int]] = None
 
-    extend_seq_lens_cpu: List[int] = None
-    logprob_start_lens_cpu: List[int] = None
+    extend_logprob_start_lens_cpu: Optional[List[int]] = None
+    extend_logprob_pruned_lens_cpu: Optional[List[int]] = None
 
     @classmethod
     def from_input_metadata(cls, input_metadata: InputMetadata):
+        return_top_logprob = any(x > 0 for x in input_metadata.top_logprobs_nums)
+        if input_metadata.forward_mode.is_extend():
+            extend_logprob_pruned_lens_cpu = [
+                extend_len - start_len
+                for extend_len, start_len in zip(
+                    input_metadata.extend_seq_lens,
+                    input_metadata.extend_logprob_start_lens_cpu,
+                )
+            ]
+        else:
+            extend_logprob_pruned_lens_cpu = None
         return cls(
             forward_mode=input_metadata.forward_mode,
-            extend_seq_lens=input_metadata.extend_seq_lens,
-            extend_start_loc=input_metadata.extend_start_loc,
-            return_logprob=input_metadata.return_logprob,
             top_logprobs_nums=input_metadata.top_logprobs_nums,
+            return_logprob=input_metadata.return_logprob,
+            return_top_logprob=return_top_logprob,
+            extend_seq_lens=input_metadata.extend_seq_lens,
             extend_seq_lens_cpu=input_metadata.extend_seq_lens_cpu,
-            logprob_start_lens_cpu=input_metadata.logprob_start_lens_cpu,
+            extend_logprob_start_lens_cpu=input_metadata.extend_logprob_start_lens_cpu,
+            extend_logprob_pruned_lens_cpu=extend_logprob_pruned_lens_cpu,
         )
 
 
@@ -82,57 +96,49 @@ class LogitsProcessor(nn.Module):
     def _get_normalized_prompt_logprobs(
         self,
         input_token_logprobs: torch.Tensor,
-        cum_start_len0: torch.Tensor,
-        cum_start_len1: torch.Tensor,
         logits_metadata: LogitsMetadata,
     ):
         logprobs_cumsum = torch.cumsum(input_token_logprobs, dim=0, dtype=torch.float32)
+        pruned_lens = torch.tensor(
+            logits_metadata.extend_logprob_pruned_lens_cpu, device="cuda"
+        )
 
-        start = logits_metadata.extend_start_loc.clone() - cum_start_len0
-        end = start + logits_metadata.extend_seq_lens - 2 - cum_start_len1
-        start.clamp_(min=0, max=input_token_logprobs.shape[0] - 1)
-        end.clamp_(min=0, max=input_token_logprobs.shape[0] - 1)
+        start = torch.zeros_like(pruned_lens)
+        start[1:] = torch.cumsum(pruned_lens[:-1], dim=0)
+        end = torch.clamp(
+            start + pruned_lens - 2, min=0, max=logprobs_cumsum.shape[0] - 1
+        )
         sum_logp = (
             logprobs_cumsum[end] - logprobs_cumsum[start] + input_token_logprobs[start]
         )
-        normalized_prompt_logprobs = sum_logp / (
-            (logits_metadata.extend_seq_lens - 1).clamp(min=1)
-        )
-
+        normalized_prompt_logprobs = sum_logp / (pruned_lens - 1).clamp(min=1)
         return normalized_prompt_logprobs
 
     @staticmethod
     def get_top_logprobs(all_logprobs: torch.Tensor, logits_metadata: LogitsMetadata):
-        if logits_metadata.forward_mode == ForwardMode.DECODE:
+        max_k = max(logits_metadata.top_logprobs_nums)
+        ret = all_logprobs.topk(max_k, dim=1)
+        values = ret.values.tolist()
+        indices = ret.indices.tolist()
+
+        if logits_metadata.forward_mode.is_decode():
             output_top_logprobs = []
-            max_k = max(logits_metadata.top_logprobs_nums)
-            ret = all_logprobs.topk(max_k, dim=1)
-            values = ret.values.tolist()
-            indices = ret.indices.tolist()
             for i, k in enumerate(logits_metadata.top_logprobs_nums):
                 output_top_logprobs.append(list(zip(values[i][:k], indices[i][:k])))
             return None, output_top_logprobs
         else:
-            # TODO: vectorize the code below
             input_top_logprobs, output_top_logprobs = [], []
-            pt = 0
-            extend_seq_lens_cpu = logits_metadata.extend_seq_lens_cpu
 
-            max_k = max(logits_metadata.top_logprobs_nums)
-            ret = all_logprobs.topk(max_k, dim=1)
-            values = ret.values.tolist()
-            indices = ret.indices.tolist()
-
-            for i, extend_seq_len in enumerate(extend_seq_lens_cpu):
-                start_len = logits_metadata.logprob_start_lens_cpu[i]
-                pruned_len = extend_seq_len - start_len
-
-                if extend_seq_len == 0:
+            pt = 0
+            for k, pruned_len in zip(
+                logits_metadata.top_logprobs_nums,
+                logits_metadata.extend_logprob_pruned_lens_cpu,
+            ):
+                if pruned_len <= 0:
                     input_top_logprobs.append([])
                     output_top_logprobs.append([])
                     continue
 
-                k = logits_metadata.top_logprobs_nums[i]
                 input_top_logprobs.append(
                     [
                         list(zip(values[pt + j][:k], indices[pt + j][:k]))
@@ -163,14 +169,11 @@ class LogitsProcessor(nn.Module):
         assert isinstance(logits_metadata, LogitsMetadata)
 
         # Get the last hidden states and last logits for the next token prediction
-        if logits_metadata.forward_mode == ForwardMode.DECODE:
+        if logits_metadata.forward_mode.is_decode():
             last_index = None
             last_hidden = hidden_states
         else:
-            last_index = (
-                torch.cumsum(logits_metadata.extend_seq_lens, dim=0, dtype=torch.long)
-                - 1
-            )
+            last_index = torch.cumsum(logits_metadata.extend_seq_lens, dim=0) - 1
             last_hidden = hidden_states[last_index]
 
         last_logits = torch.matmul(last_hidden, weight.T)
@@ -194,21 +197,15 @@ class LogitsProcessor(nn.Module):
                 output_top_logprobs=None,
             )
         else:
-            # When logprob is requested, compute the logits for all tokens.
-            if logits_metadata.forward_mode == ForwardMode.DECODE:
-                last_logprobs = torch.nn.functional.log_softmax(last_logits, dim=-1)
+            last_logprobs = torch.nn.functional.log_softmax(last_logits, dim=-1)
 
-                # Get the logprob of top-k tokens
-                return_top_logprob = any(
-                    x > 0 for x in logits_metadata.top_logprobs_nums
-                )
-                if return_top_logprob:
+            if logits_metadata.forward_mode.is_decode():
+                if logits_metadata.return_top_logprob:
                     output_top_logprobs = self.get_top_logprobs(
                         last_logprobs, logits_metadata
                     )[1]
                 else:
                     output_top_logprobs = None
-
                 return LogitsProcessorOutput(
                     next_token_logits=last_logits,
                     next_token_logprobs=last_logprobs,
@@ -218,22 +215,18 @@ class LogitsProcessor(nn.Module):
                     output_top_logprobs=output_top_logprobs,
                 )
             else:
+                # Slice the requested tokens to compute logprob
                 pt, states, pruned_input_ids = 0, [], []
-                for i, extend_len in enumerate(logits_metadata.extend_seq_lens_cpu):
-                    start_len = logits_metadata.logprob_start_lens_cpu[i]
+                for start_len, extend_len in zip(
+                    logits_metadata.extend_logprob_start_lens_cpu,
+                    logits_metadata.extend_seq_lens_cpu,
+                ):
                     states.append(hidden_states[pt + start_len : pt + extend_len])
                     pruned_input_ids.append(input_ids[pt + start_len : pt + extend_len])
                     pt += extend_len
 
+                # Compute the logits and logprobs for all required tokens
                 states = torch.cat(states, dim=0)
-                pruned_input_ids = torch.cat(pruned_input_ids, dim=0)
-
-                cum_start_len1 = torch.tensor(
-                    logits_metadata.logprob_start_lens_cpu, device="cuda"
-                ).cumsum(0)
-                cum_start_len0 = torch.zeros_like(cum_start_len1)
-                cum_start_len0[1:] = cum_start_len1[:-1]
-
                 all_logits = torch.matmul(states, weight.T)
                 if self.do_tensor_parallel_all_gather:
                     all_logits = tensor_model_parallel_all_gather(all_logits)
@@ -249,35 +242,29 @@ class LogitsProcessor(nn.Module):
                 all_logprobs[:] = torch.nn.functional.log_softmax(all_logprobs, dim=-1)
 
                 # Get the logprob of top-k tokens
-                return_top_logprob = any(
-                    x > 0 for x in logits_metadata.top_logprobs_nums
-                )
-                if return_top_logprob:
+                if logits_metadata.return_top_logprob:
                     input_top_logprobs, output_top_logprobs = self.get_top_logprobs(
                         all_logprobs, logits_metadata
                     )
                 else:
                     input_top_logprobs = output_top_logprobs = None
 
-                last_logprobs = all_logprobs[last_index - cum_start_len1]
-
-                # Compute the logprobs and normalized logprobs for the prefill tokens.
-                # Note that we pad a zero at the end of each sequence for easy computation.
+                # Compute the normalized logprobs for the requested tokens.
+                # Note that we pad a zero at the end for easy batching.
                 input_token_logprobs = all_logprobs[
                     torch.arange(all_logprobs.shape[0], device="cuda"),
-                    torch.cat([pruned_input_ids[1:], torch.tensor([0], device="cuda")]),
+                    torch.cat(
+                        [
+                            torch.cat(pruned_input_ids)[1:],
+                            torch.tensor([0], device="cuda"),
+                        ]
+                    ),
                 ]
-
                 normalized_prompt_logprobs = self._get_normalized_prompt_logprobs(
                     input_token_logprobs,
-                    cum_start_len0,
-                    cum_start_len1,
                     logits_metadata,
                 )
 
-                # Remove the last token logprob for the prefill tokens.
-                input_token_logprobs = input_token_logprobs[:-1]
-
                 return LogitsProcessorOutput(
                     next_token_logits=last_logits,
                     next_token_logprobs=last_logprobs,

sglang/srt/layers/radix_attention.py

@@ -15,20 +15,16 @@ limitations under the License.
 
 """Radix attention."""
 
-from typing import Optional
-
-import torch
-from flashinfer.cascade import merge_state
 from torch import nn
 
-from sglang.global_config import global_config
-from sglang.srt.layers.decode_attention import decode_attention_fwd
-from sglang.srt.layers.extend_attention import extend_attention_fwd
-from sglang.srt.model_executor.forward_batch_info import ForwardMode, InputMetadata
-from sglang.srt.model_executor.model_runner import global_server_args_dict
+from sglang.srt.model_executor.forward_batch_info import InputMetadata
 
 
 class RadixAttention(nn.Module):
+    """
+    The attention layer implementation.
+    """
+
     def __init__(
         self,
         num_heads: int,
@@ -36,8 +32,8 @@ class RadixAttention(nn.Module):
         scaling: float,
         num_kv_heads: int,
         layer_id: int,
-        sliding_window_size: Optional[int] = None,
-        logit_cap: int = -1,
+        sliding_window_size: int = -1,
+        logit_cap: float = 0.0,
         v_head_dim: int = -1,
     ):
         super().__init__()
@@ -49,160 +45,14 @@ class RadixAttention(nn.Module):
         self.v_head_dim = v_head_dim if v_head_dim != -1 else head_dim
         self.scaling = scaling
         self.layer_id = layer_id
-        self.sliding_window_size = sliding_window_size if sliding_window_size else -1
-
-        if (
-            not global_server_args_dict.get("disable_flashinfer", False)
-            and self.qk_head_dim == self.v_head_dim
-        ):
-            self.extend_forward = self.extend_forward_flashinfer
-            self.decode_forward = self.decode_forward_flashinfer
-        else:
-            self.extend_forward = self.extend_forward_triton
-            self.decode_forward = self.decode_forward_triton
-
-        self.logit_cap = logit_cap if logit_cap is not None and logit_cap > 0 else 0
-
-    def extend_forward_triton(self, q, k, v, input_metadata: InputMetadata):
-        if self.qk_head_dim != self.v_head_dim:
-            o = q.new_empty((q.shape[0], self.tp_q_head_num * self.v_head_dim))
-        else:
-            o = torch.empty_like(q)
-
-        self.store_kv_cache(k, v, input_metadata)
-        extend_attention_fwd(
-            q.view(-1, self.tp_q_head_num, self.qk_head_dim),
-            k.contiguous(),
-            v.contiguous(),
-            o.view(-1, self.tp_q_head_num, self.v_head_dim),
-            input_metadata.token_to_kv_pool.get_key_buffer(self.layer_id),
-            input_metadata.token_to_kv_pool.get_value_buffer(self.layer_id),
-            input_metadata.req_to_token_pool.req_to_token,
-            input_metadata.req_pool_indices,
-            input_metadata.triton_start_loc,
-            input_metadata.seq_lens,
-            input_metadata.triton_prefix_lens,
-            input_metadata.extend_start_loc,
-            input_metadata.extend_seq_lens,
-            input_metadata.triton_max_seq_len,
-            input_metadata.triton_max_extend_len,
-            sm_scale=self.scaling,
-            logit_cap=self.logit_cap,
-        )
-
-        return o
-
-    def decode_forward_triton(self, q, k, v, input_metadata: InputMetadata):
-        if self.qk_head_dim != self.v_head_dim:
-            o = q.new_empty((q.shape[0], self.tp_q_head_num * self.v_head_dim))
-        else:
-            o = torch.empty_like(q)
-        self.store_kv_cache(k, v, input_metadata)
-
-        decode_attention_fwd(
-            q.view(-1, self.tp_q_head_num, self.qk_head_dim),
-            input_metadata.token_to_kv_pool.get_key_buffer(self.layer_id),
-            input_metadata.token_to_kv_pool.get_value_buffer(self.layer_id),
-            o.view(-1, self.tp_q_head_num, self.v_head_dim),
-            input_metadata.req_to_token_pool.req_to_token,
-            input_metadata.req_pool_indices,
-            input_metadata.triton_start_loc,
-            input_metadata.seq_lens,
-            input_metadata.triton_max_seq_len,
-            input_metadata.total_num_tokens,
-            sm_scale=self.scaling,
-            logit_cap=self.logit_cap,
-        )
-
-        return o
-
-    def extend_forward_flashinfer(self, q, k, v, input_metadata: InputMetadata):
-        # using two wrappers is unnecessary in the current PR, but are prepared for future PRs
-        prefill_wrapper_paged = input_metadata.flashinfer_prefill_wrapper_paged
-        if self.sliding_window_size != -1:
-            prefill_wrapper_paged = prefill_wrapper_paged[0]
-        else:
-            if isinstance(prefill_wrapper_paged, list):
-                prefill_wrapper_paged = prefill_wrapper_paged[1]
-
-        if not input_metadata.flashinfer_use_ragged:
-            if k is not None:
-                assert v is not None
-                self.store_kv_cache(k, v, input_metadata)
-
-            o = prefill_wrapper_paged.forward(
-                q.contiguous().view(-1, self.tp_q_head_num, self.head_dim),
-                input_metadata.token_to_kv_pool.get_kv_buffer(self.layer_id),
-                causal=True,
-                sm_scale=self.scaling,
-                window_left=self.sliding_window_size,
-                logits_soft_cap=self.logit_cap,
-            )
-        else:
-            o1, s1 = (
-                input_metadata.flashinfer_prefill_wrapper_ragged.forward_return_lse(
-                    q.contiguous().view(-1, self.tp_q_head_num, self.head_dim),
-                    k.contiguous().view(-1, self.tp_k_head_num, self.head_dim),
-                    v.contiguous().view(-1, self.tp_v_head_num, self.head_dim),
-                    causal=True,
-                    sm_scale=self.scaling,
-                    logits_soft_cap=self.logit_cap,
-                )
-            )
-
-            if input_metadata.extend_no_prefix:
-                o = o1
-            else:
-                o2, s2 = prefill_wrapper_paged.forward_return_lse(
-                    q.contiguous().view(-1, self.tp_q_head_num, self.head_dim),
-                    input_metadata.token_to_kv_pool.get_kv_buffer(self.layer_id),
-                    causal=False,
-                    sm_scale=self.scaling,
-                    logits_soft_cap=self.logit_cap,
-                )
-
-                o, _ = merge_state(o1, s1, o2, s2)
-
-            self.store_kv_cache(k, v, input_metadata)
-
-            if input_metadata.total_num_tokens >= global_config.layer_sync_threshold:
-                torch.cuda.synchronize()
-
-        return o.view(-1, self.tp_q_head_num * self.head_dim)
-
-    def decode_forward_flashinfer(self, q, k, v, input_metadata: InputMetadata):
-        decode_wrapper = input_metadata.flashinfer_decode_wrapper
-        if self.sliding_window_size != -1:
-            decode_wrapper = decode_wrapper[0]
-        else:
-            if isinstance(decode_wrapper, list):
-                decode_wrapper = decode_wrapper[1]
-
-        if k is not None:
-            assert v is not None
-            self.store_kv_cache(k, v, input_metadata)
-
-        o = decode_wrapper.forward(
-            q.contiguous().view(-1, self.tp_q_head_num, self.head_dim),
-            input_metadata.token_to_kv_pool.get_kv_buffer(self.layer_id),
-            sm_scale=self.scaling,
-            logits_soft_cap=self.logit_cap,
-        )
-
-        return o.view(-1, self.tp_q_head_num * self.head_dim)
+        self.logit_cap = logit_cap
+        self.sliding_window_size = sliding_window_size or -1
 
     def forward(self, q, k, v, input_metadata: InputMetadata):
         if k is not None:
+            # For cross-layer sharing, kv can be None
             assert v is not None
             k = k.view(-1, self.tp_k_head_num, self.qk_head_dim)
             v = v.view(-1, self.tp_v_head_num, self.v_head_dim)
 
-        if input_metadata.forward_mode == ForwardMode.EXTEND:
-            return self.extend_forward(q, k, v, input_metadata)
-        elif input_metadata.forward_mode == ForwardMode.DECODE:
-            return self.decode_forward(q, k, v, input_metadata)
-
-    def store_kv_cache(self, cache_k, cache_v, input_metadata: InputMetadata):
-        input_metadata.token_to_kv_pool.set_kv_buffer(
-            self.layer_id, input_metadata.out_cache_loc, cache_k, cache_v
-        )
+        return input_metadata.attn_backend.forward(q, k, v, self, input_metadata)