sglang 0.3.2__py3-none-any.whl → 0.3.3__py3-none-any.whl

sglang/srt/layers/attention/__init__.py

@@ -0,0 +1,49 @@
+from abc import ABC, abstractmethod
+
+from torch import nn
+
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+
+
+class AttentionBackend(ABC):
+    """The base class of attention backends"""
+
+    @abstractmethod
+    def init_forward_metadata(self, forward_batch: ForwardBatch):
+        """Init the metadata for a forward pass."""
+        raise NotImplementedError()
+
+    def init_cuda_graph_state(self, max_bs: int):
+        """Init the global shared states for cuda graph."""
+        raise NotImplementedError()
+
+    def init_forward_metadata_capture_cuda_graph(
+        self, bs: int, req_pool_indices, seq_lens
+    ):
+        """Init the metadata for a forward pass for capturing a cuda graph."""
+        raise NotImplementedError()
+
+    def init_forward_metadata_replay_cuda_graph(
+        self, bs: int, req_pool_indices, seq_lens
+    ):
+        """Init the metadata for a forward pass for replying a cuda graph."""
+        raise NotImplementedError()
+
+    def get_cuda_graph_seq_len_fill_value(self):
+        """Get the fill value for padded seq lens. Typically, it is 0 or 1."""
+        raise NotImplementedError()
+
+    def forward(self, q, k, v, layer: nn.Module, forward_batch: ForwardBatch):
+        """Run forward on an attention layer."""
+        if forward_batch.forward_mode.is_decode():
+            return self.forward_decode(q, k, v, layer, forward_batch)
+        else:
+            return self.forward_extend(q, k, v, layer, forward_batch)
+
+    def forward_decode(self, q, k, v, layer: nn.Module, forward_batch: ForwardBatch):
+        """Run a forward for decode."""
+        raise NotImplementedError()
+
+    def forward_extend(self, q, k, v, layer: nn.Module, forward_batch: ForwardBatch):
+        """Run a forward for extend."""
+        raise NotImplementedError()

sglang/srt/layers/attention/flashinfer_backend.py

@@ -0,0 +1,277 @@
+from __future__ import annotations
+
+"""
+Support different attention backends.
+Now there are two backends: FlashInfer and Triton.
+FlashInfer is faster and Triton is easier to customize.
+Each backend supports two operators: extend (i.e. prefill with cached prefix) and decode.
+"""
+
+from typing import TYPE_CHECKING
+
+import torch
+import torch.nn as nn
+
+from sglang.global_config import global_config
+from sglang.srt.layers.attention import AttentionBackend
+from sglang.srt.layers.attention.flashinfer_utils import (
+    WrapperDispatch,
+    update_flashinfer_indices,
+)
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
+from sglang.srt.utils import is_flashinfer_available
+
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.model_runner import ModelRunner
+
+if is_flashinfer_available():
+    from flashinfer import (
+        BatchDecodeWithPagedKVCacheWrapper,
+        BatchPrefillWithPagedKVCacheWrapper,
+        BatchPrefillWithRaggedKVCacheWrapper,
+    )
+    from flashinfer.cascade import merge_state
+    from flashinfer.decode import _grouped_size_compiled_for_decode_kernels
+
+
+class FlashInferAttnBackend(AttentionBackend):
+    """Flashinfer attention kernels."""
+
+    def __init__(self, model_runner: ModelRunner):
+        super().__init__()
+        self.model_runner = model_runner
+
+        if not _grouped_size_compiled_for_decode_kernels(
+            model_runner.model_config.num_attention_heads // model_runner.tp_size,
+            model_runner.model_config.get_num_kv_heads(model_runner.tp_size),
+        ):
+            self.decode_use_tensor_cores = True
+        else:
+            self.decode_use_tensor_cores = False
+
+        self.workspace_buffer = torch.empty(
+            global_config.flashinfer_workspace_size,
+            dtype=torch.uint8,
+            device="cuda",
+        )
+
+        assert not (
+            model_runner.sliding_window_size is not None
+            and model_runner.has_cross_attention
+        ), "Sliding window and cross attention are not supported together"
+
+        self.num_wrappers = 1
+        self.dispatch_reason = None
+        if model_runner.sliding_window_size is not None:
+            self.num_wrappers = 2
+            self.dispatch_reason = WrapperDispatch.SLIDING_WINDOW
+        elif model_runner.has_cross_attention:
+            self.num_wrappers = 2
+            self.dispatch_reason = WrapperDispatch.CROSS_ATTENTION
+
+        # NOTE: we do not use ragged attention when there are multiple wrappers
+        self.prefill_wrapper_ragged = (
+            BatchPrefillWithRaggedKVCacheWrapper(self.workspace_buffer, "NHD")
+            if self.num_wrappers == 1
+            else None
+        )
+
+        # Two wrappers: one for sliding window attention and one for full attention.
+        # Using two wrappers is unnecessary in the current PR, but are prepared for future PRs
+        self.prefill_wrappers_paged = []
+        self.decode_wrappers = []
+        for _ in range(self.num_wrappers):
+            self.prefill_wrappers_paged.append(
+                BatchPrefillWithPagedKVCacheWrapper(self.workspace_buffer, "NHD")
+            )
+            self.decode_wrappers.append(
+                BatchDecodeWithPagedKVCacheWrapper(
+                    self.workspace_buffer,
+                    "NHD",
+                    use_tensor_cores=self.decode_use_tensor_cores,
+                )
+            )
+
+        self.forward_metadata = None
+        self.cuda_graph_metadata = {}
+
+    def _get_wrapper_idx(self, layer: nn.Module):
+        if self.num_wrappers == 1:
+            return 0
+
+        if self.dispatch_reason == WrapperDispatch.SLIDING_WINDOW:
+            return layer.sliding_window_size == -1
+        if self.dispatch_reason == WrapperDispatch.CROSS_ATTENTION:
+            return layer.is_cross_attention
+
+        raise ValueError(f"Unknown dispatch reason: {self.dispatch_reason}")
+
+    def init_forward_metadata(self, forward_batch: ForwardBatch):
+        if forward_batch.forward_mode.is_decode():
+            prefix_lens = None
+            use_ragged = False
+            extend_no_prefix = False
+            total_num_tokens = None
+        else:
+            prefix_lens = forward_batch.extend_prefix_lens
+
+            # Some heuristics to check whether to use ragged forward
+            use_ragged = False
+            if (
+                torch.sum(forward_batch.seq_lens).item() >= 4096
+                and self.num_wrappers == 1
+            ):
+                use_ragged = True
+
+            total_num_tokens = torch.sum(forward_batch.seq_lens).item()
+            extend_no_prefix = not torch.any(forward_batch.extend_prefix_lens).item()
+
+        update_flashinfer_indices(
+            forward_batch.forward_mode,
+            self.model_runner,
+            forward_batch.req_pool_indices,
+            forward_batch.seq_lens,
+            prefix_lens,
+            use_ragged=use_ragged,
+        )
+
+        self.forward_metadata = (
+            use_ragged,
+            extend_no_prefix,
+            total_num_tokens,
+            self.decode_wrappers,
+        )
+
+    def init_cuda_graph_state(self, max_bs: int):
+        self.cuda_graph_kv_indptr = torch.zeros(
+            (max_bs + 1,), dtype=torch.int32, device="cuda"
+        )
+        self.cuda_graph_kv_indices = torch.zeros(
+            (max_bs * self.model_runner.model_config.context_len,),
+            dtype=torch.int32,
+            device="cuda",
+        )
+        self.cuda_graph_kv_last_page_len = torch.ones(
+            (max_bs,), dtype=torch.int32, device="cuda"
+        )
+
+        # NOTE: the buffers are always in the form of list
+        self.cuda_graph_kv_indptr = [self.cuda_graph_kv_indptr] + [
+            self.cuda_graph_kv_indptr.clone() for _ in range(self.num_wrappers - 1)
+        ]
+        self.cuda_graph_kv_indices = [self.cuda_graph_kv_indices] + [
+            self.cuda_graph_kv_indices.clone() for _ in range(self.num_wrappers - 1)
+        ]
+
+    def init_forward_metadata_capture_cuda_graph(
+        self, bs: int, req_pool_indices, seq_lens
+    ):
+        decode_wrappers = []
+        for i in range(self.num_wrappers):
+            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.cuda_graph_kv_indptr[i][: bs + 1],
+                    paged_kv_indices_buffer=self.cuda_graph_kv_indices[i],
+                    paged_kv_last_page_len_buffer=self.cuda_graph_kv_last_page_len[:bs],
+                )
+            )
+
+        update_flashinfer_indices(
+            ForwardMode.DECODE,
+            self.model_runner,
+            req_pool_indices,
+            seq_lens,
+            None,
+            decode_wrappers,
+        )
+
+        self.cuda_graph_metadata[bs] = decode_wrappers
+
+        self.forward_metadata = (False, False, None, decode_wrappers)
+
+    def init_forward_metadata_replay_cuda_graph(
+        self, bs: int, req_pool_indices, seq_lens
+    ):
+        update_flashinfer_indices(
+            ForwardMode.DECODE,
+            self.model_runner,
+            req_pool_indices[:bs],
+            seq_lens[:bs],
+            None,
+            self.cuda_graph_metadata[bs],
+        )
+
+    def get_cuda_graph_seq_len_fill_value(self):
+        return 0
+
+    def forward_extend(self, q, k, v, layer: nn.Module, forward_batch: ForwardBatch):
+        prefill_wrapper_paged = self.prefill_wrappers_paged[
+            self._get_wrapper_idx(layer)
+        ]
+
+        use_ragged, extend_no_prefix, _, _ = self.forward_metadata
+
+        if not use_ragged:
+            if k is not None:
+                assert v is not None
+                forward_batch.token_to_kv_pool.set_kv_buffer(
+                    layer.layer_id, forward_batch.out_cache_loc, k, v
+                )
+            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=True,
+                sm_scale=layer.scaling,
+                window_left=layer.sliding_window_size,
+                logits_soft_cap=layer.logit_cap,
+            )
+        else:
+            o1, s1 = self.prefill_wrapper_ragged.forward_return_lse(
+                q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim),
+                k.contiguous().view(-1, layer.tp_k_head_num, layer.head_dim),
+                v.contiguous().view(-1, layer.tp_v_head_num, layer.head_dim),
+                causal=True,
+                sm_scale=layer.scaling,
+                logits_soft_cap=layer.logit_cap,
+            )
+
+            if 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)
+
+            forward_batch.token_to_kv_pool.set_kv_buffer(
+                layer.layer_id, forward_batch.out_cache_loc, k, v
+            )
+
+        return o.view(-1, layer.tp_q_head_num * layer.head_dim)
+
+    def forward_decode(self, q, k, v, layer: nn.Module, forward_batch: ForwardBatch):
+        decode_wrapper = self.forward_metadata[-1][self._get_wrapper_idx(layer)]
+
+        if k is not None:
+            assert v is not None
+            forward_batch.token_to_kv_pool.set_kv_buffer(
+                layer.layer_id, forward_batch.out_cache_loc, k, v
+            )
+
+        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,
+        )
+
+        return o.view(-1, layer.tp_q_head_num * layer.head_dim)

sglang/srt/layers/{flashinfer_utils.py → attention/flashinfer_utils.py}

@@ -1,8 +1,15 @@
+from enum import Enum, auto
+
 import torch
 import triton
 import triton.language as tl
 
 
+class WrapperDispatch(Enum):
+    SLIDING_WINDOW = auto()
+    CROSS_ATTENTION = auto()
+
+
 @triton.jit
 def create_flashinfer_kv_indices_triton(
     req_to_token_ptr,  # [max_batch, max_context_len]
@@ -47,7 +54,7 @@ class FlashinferUpdater:
         req_pool_indices,
         seq_lens,
         prefix_lens,
-        decode_wrapper=None,
+        decode_wrappers=None,
         use_ragged=False,
     ):
         self.forward_mode = forward_mode
@@ -66,82 +73,22 @@ class FlashinferUpdater:
         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.decode_wrappers = (
+            decode_wrappers or self.model_runner.attn_backend.decode_wrappers
         )
         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.prefill_wrappers_paged = (
+            self.model_runner.attn_backend.prefill_wrappers_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):
+        assert not isinstance(decode_wrapper, list)
         decode_wrapper.end_forward()
         decode_wrapper.begin_forward(
             self.kv_indptr,
@@ -156,6 +103,9 @@ class FlashinferUpdater:
         )
 
     def _update_extend_indices(self, ragged_wrapper, paged_wrapper):
+        assert not isinstance(paged_wrapper, list)
+        assert not isinstance(ragged_wrapper, list)
+
         # extend part
         qo_indptr = torch.zeros(
             (self.batch_size + 1,), dtype=torch.int32, device="cuda"
@@ -185,28 +135,75 @@ class FlashinferUpdater:
             1,
         )
 
-    def update_indices_no_sliding_window(self):
-        self._init_indices_no_sliding_window()
+    def _get_indices(self, dispatch_reason: WrapperDispatch = None, wrapper_id=0):
+        if dispatch_reason is None:
+            if self.use_ragged:
+                paged_kernel_lens = self.prefix_lens
+            else:
+                paged_kernel_lens = self.seq_lens
+            self.kv_start_idx = None
+        elif dispatch_reason == WrapperDispatch.SLIDING_WINDOW:
+            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
+            self.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,
+            self.kv_start_idx,
+            self.kv_indices,
+            self.model_runner.req_to_token_pool.req_to_token.size(1),
+        )
+
+    def _update_indicess_single_wrapper(self):
+        self._get_indices()
 
         if self.forward_mode.is_decode():
-            self._update_decode_indices(self.decode_wrapper)
+            self._update_decode_indices(self.decode_wrappers[0])
         else:
             self._update_extend_indices(
                 self.prefill_wrapper_ragged,
-                self.prefill_wrapper_paged,
+                self.prefill_wrappers_paged[0],
             )
 
-    def update_indices_sliding_window(self):
-        assert self.use_ragged is False
+    def _update_indices_cross_attention(self):
+        pass
 
+    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)
+            self._get_indices(WrapperDispatch.SLIDING_WINDOW, wrapper_id)
             if self.forward_mode.is_decode():
-                self._update_decode_indices(self.decode_wrapper[wrapper_id])
+                self._update_decode_indices(self.decode_wrappers[wrapper_id])
             else:
                 self._update_extend_indices(
                     None,
-                    self.prefill_wrapper_paged[wrapper_id],
+                    self.prefill_wrappers_paged[wrapper_id],
                 )
 
 
@@ -216,7 +213,7 @@ def update_flashinfer_indices(
     req_pool_indices,
     seq_lens,
     prefix_lens,
-    decode_wrapper=None,
+    decode_wrappers=None,
     use_ragged=False,
 ):
     updater = FlashinferUpdater(
@@ -225,11 +222,16 @@ def update_flashinfer_indices(
         req_pool_indices,
         seq_lens,
         prefix_lens,
-        decode_wrapper,
+        decode_wrappers,
         use_ragged,
     )
 
-    if model_runner.sliding_window_size is None:
-        updater.update_indices_no_sliding_window()
+    dispatch_reason = model_runner.attn_backend.dispatch_reason
+
+    if dispatch_reason == WrapperDispatch.SLIDING_WINDOW:
+        updater._update_indices_sliding_window()
+    elif dispatch_reason == WrapperDispatch.CROSS_ATTENTION:
+        updater._update_indices_cross_attention()
     else:
-        updater.update_indices_sliding_window()
+        assert model_runner.attn_backend.num_wrappers == 1
+        updater._update_indicess_single_wrapper()