chore: import upstream snapshot with attribution

python/sglang/srt/mem_cache/sparsity/backend/__init__.py

@@ -0,0 +1,7 @@
+from sglang.srt.mem_cache.sparsity.backend.backend_adaptor import (
+    BackendAdaptor,
+    DSABackendAdaptor,
+    FlashAttentionAdaptor,
+)
+
+__all__ = ["BackendAdaptor", "FlashAttentionAdaptor", "DSABackendAdaptor"]

python/sglang/srt/mem_cache/sparsity/backend/backend_adaptor.py

@@ -0,0 +1,176 @@
+import logging
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any, Optional
+
+import torch
+
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+
+logger = logging.getLogger(__name__)
+
+
+class BackendAdaptor(ABC):
+    """Base class for attention backend adaptors."""
+
+    def __init__(self, device: torch.device):
+        self.device = device
+        self._original_metadata = None
+
+    def save_original_metadata(self, metadata: Any) -> None:
+        """Save original metadata in the beginning of the forward pass."""
+        pass
+
+    @abstractmethod
+    def adapt_for_attn_metadata(
+        self,
+        selected_indices: torch.Tensor,
+        valid_lengths: torch.Tensor,
+        sparse_mask: torch.Tensor,
+        current_metadata: Any,
+        forward_batch: "ForwardBatch",
+        req_to_token: torch.Tensor,
+        page_size: int,
+        layer_id: int,
+        **kwargs,
+    ) -> Any:
+        """
+        Adapt attention metadata for sparse KVCache access.
+
+        Transforms sparse retrieval results (logical indices of important KV pages/tokens)
+        into backend-specific attention metadata format.
+
+        Returns:
+            Modified attention metadata compatible with the backend
+        """
+        pass
+
+
+class DSABackendAdaptor(BackendAdaptor):
+    """Adaptor for DSA (DeepSeek Sparse Attention) backend."""
+
+    def __init__(
+        self,
+        device: torch.device,
+        req_to_token_pool,
+    ):
+        super().__init__(device)
+        self.req_to_token_pool = req_to_token_pool
+
+    def adapt_for_attn_metadata(
+        self,
+        selected_indices: torch.Tensor,
+        valid_lengths: torch.Tensor,
+        sparse_mask: torch.Tensor,
+        current_metadata: Any,
+        forward_batch: "ForwardBatch",
+        req_to_token: torch.Tensor,
+        page_size: int,
+        layer_id: int,
+        **kwargs,
+    ) -> Optional[torch.Tensor]:
+        """
+        Transform logical page indices to physical device indices for DSA backend.
+        """
+        # TODO: Implement DSA backend adaptor logic
+        pass
+
+
+class FlashAttentionAdaptor(BackendAdaptor):
+    """Adaptor for FlashAttention backend."""
+
+    def save_original_metadata(self, metadata: Any) -> None:
+        self._original_metadata = {
+            "page_table": metadata.page_table.clone(),
+            "cache_seqlens_int32": metadata.cache_seqlens_int32.clone(),
+            "cu_seqlens_k": metadata.cu_seqlens_k.clone(),
+            "max_seq_len_k": metadata.max_seq_len_k,
+        }
+
+    def adapt_for_attn_metadata(
+        self,
+        selected_indices: torch.Tensor,
+        valid_lengths: torch.Tensor,
+        sparse_mask: torch.Tensor,
+        current_metadata: Any,
+        forward_batch: "ForwardBatch",
+        req_to_token: torch.Tensor,
+        page_size: int,
+        layer_id: int,
+        **kwargs,
+    ) -> Any:
+        """
+        Adapt FlashAttention metadata for sparse KVCache access.
+
+        Modifies page_table, cache_seqlens, and related metadata to redirect
+        FlashAttention to only process selected sparse pages.
+
+        # TODO: Optimize performance
+        """
+        if self._original_metadata is None:
+            return current_metadata
+
+        if not sparse_mask.any():
+            return current_metadata
+
+        current_metadata.page_table.copy_(self._original_metadata["page_table"])
+        current_metadata.cache_seqlens_int32.copy_(
+            self._original_metadata["cache_seqlens_int32"]
+        )
+
+        physical_pages = self._logical_to_physical_pages_batch(
+            selected_indices,
+            forward_batch.req_pool_indices,
+            req_to_token,
+            page_size,
+        )
+
+        max_selected = physical_pages.shape[1]
+        valid_mask = torch.arange(max_selected, device=physical_pages.device).unsqueeze(
+            0
+        ) < valid_lengths.unsqueeze(1)
+        update_mask = sparse_mask.unsqueeze(1) & valid_mask
+
+        current_metadata.page_table[:, :max_selected] = torch.where(
+            update_mask, physical_pages, current_metadata.page_table[:, :max_selected]
+        )
+
+        seq_lens = forward_batch.seq_lens
+        positions_in_page = (seq_lens - 1) % page_size
+        diff = page_size - positions_in_page - 1
+        sparse_seq_lens = (valid_lengths * page_size - diff).to(torch.int32)
+
+        current_metadata.cache_seqlens_int32 = torch.where(
+            sparse_mask, sparse_seq_lens, self._original_metadata["cache_seqlens_int32"]
+        )
+
+        current_metadata.cu_seqlens_k = torch.nn.functional.pad(
+            torch.cumsum(
+                current_metadata.cache_seqlens_int32, dim=0, dtype=torch.int32
+            ),
+            (1, 0),
+        )
+        current_metadata.max_seq_len_k = int(current_metadata.cache_seqlens_int32.max())
+        return current_metadata
+
+    def _logical_to_physical_pages_batch(
+        self,
+        logical_pages: torch.Tensor,
+        req_pool_indices: torch.Tensor,
+        req_to_token: torch.Tensor,
+        page_size: int,
+    ) -> torch.Tensor:
+        bs, max_pages = logical_pages.shape
+
+        page_starts = logical_pages * page_size
+        page_starts_clamped = page_starts.clamp(min=0)
+
+        req_indices_expanded = req_pool_indices.unsqueeze(1).expand(-1, max_pages)
+        first_tokens = req_to_token[req_indices_expanded, page_starts_clamped]
+
+        physical_pages = first_tokens // page_size
+        physical_pages = torch.where(
+            logical_pages >= 0, physical_pages, torch.zeros_like(physical_pages)
+        )
+
+        return physical_pages.to(torch.int32)