chore: import upstream snapshot with attribution

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

@@ -0,0 +1,29 @@
+from sglang.srt.mem_cache.sparsity.algorithms import (
+    BaseSparseAlgorithm,
+    BaseSparseAlgorithmImpl,
+    DeepSeekDSAAlgorithm,
+    QuestAlgorithm,
+)
+from sglang.srt.mem_cache.sparsity.backend import BackendAdaptor, FlashAttentionAdaptor
+from sglang.srt.mem_cache.sparsity.core import SparseConfig, SparseCoordinator
+from sglang.srt.mem_cache.sparsity.factory import (
+    create_sparse_coordinator,
+    get_sparse_coordinator,
+    parse_hisparse_config,
+    register_sparse_coordinator,
+)
+
+__all__ = [
+    "BaseSparseAlgorithm",
+    "BaseSparseAlgorithmImpl",
+    "QuestAlgorithm",
+    "DeepSeekDSAAlgorithm",
+    "BackendAdaptor",
+    "FlashAttentionAdaptor",
+    "SparseConfig",
+    "SparseCoordinator",
+    "create_sparse_coordinator",
+    "get_sparse_coordinator",
+    "parse_hisparse_config",
+    "register_sparse_coordinator",
+]

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

@@ -0,0 +1,13 @@
+from sglang.srt.mem_cache.sparsity.algorithms.base_algorithm import (
+    BaseSparseAlgorithm,
+    BaseSparseAlgorithmImpl,
+)
+from sglang.srt.mem_cache.sparsity.algorithms.deepseek_dsa import DeepSeekDSAAlgorithm
+from sglang.srt.mem_cache.sparsity.algorithms.quest_algorithm import QuestAlgorithm
+
+__all__ = [
+    "BaseSparseAlgorithm",
+    "BaseSparseAlgorithmImpl",
+    "DeepSeekDSAAlgorithm",
+    "QuestAlgorithm",
+]

python/sglang/srt/mem_cache/sparsity/algorithms/base_algorithm.py

@@ -0,0 +1,383 @@
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING
+
+import torch
+
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+
+
+class BaseSparseAlgorithm(ABC):
+    """
+    Abstract base class for sparse attention algorithms.
+
+    This class provides a unified interface for implementing various retrievable KVCache
+    compression algorithms. Token-wise sparsity is treated as page-wise with page_size=1.
+
+    References:
+        - ChunkKV: https://arxiv.org/abs/2502.00299
+        - Quest: https://arxiv.org/pdf/2406.10774
+        - PQCache: https://arxiv.org/abs/2407.12820
+        - SnapKV: https://arxiv.org/pdf/2404.14469
+        - Look-ahead QCache: https://arxiv.org/pdf/2505.20334
+        - and more...
+    """
+
+    def __init__(self, config, device: torch.device, **kwargs):
+        self.config = config
+        self.device = device
+        self.req_to_token_pool = None
+        self.states = None
+
+    def initialize_representation_pool(
+        self,
+        start_layer: int,
+        end_layer: int,
+        token_to_kv_pool,
+        req_to_token_pool,
+        states,
+    ):
+        """
+        Initialize algorithm-specific representation pool and set context.
+
+        Called once during SparseCoordinator initialization. Algorithms allocate
+        their own representation tensors and store references to context.
+
+        Algorithm-specific implementations:
+            - ChunkKV: Allocate chunk scores [num_chunks, 1] for tracking semantic chunk importance
+            - Quest: Allocate page representations [num_pages, repr_dim] via key pooling
+            - PQCache: Allocate centroids [n_subvec, n_centroids, subvec_dim] and token codes [num_tokens, n_subvec]
+            - SnapKV: Allocate voting scores [num_tokens] and selected positions mask for retention strategy
+            - Look-ahead QCache: Allocate importance scores [num_tokens], eviction mask, and optional pseudo query cache [cache_size, hidden_dim]
+        """
+        pass
+
+    def construct_representations(
+        self,
+        layer_id: int,
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        k_buffer: torch.Tensor,
+        forward_batch: "ForwardBatch",
+    ):
+        """
+        Construct initial representations during prefill phase.
+
+        Called at every layer during forward pass. Algorithm internally decides
+        whether to perform construction.
+        Typically only constructs once per request during prefill/extend phase.
+
+        Algorithm-specific implementations:
+            - ChunkKV: Compute chunk importance scores via aggregated key L2 norms within semantic chunks
+            - Quest: Compute page representations via mean pooling of keys within each page
+            - PQCache: Run K-means clustering to generate centroids and assign each token to nearest centroid
+            - SnapKV: Select observation window (recent tokens), compute attention weights, aggregate via voting to identify important prefix positions, apply 1D pooling to preserve context
+            - Look-ahead QCache: Generate pseudo lookahead query (e.g., mean of last k queries), compute KV importance scores, mark low-importance KVs for eviction
+        """
+        pass
+
+    def update_representations(
+        self,
+        layer_id: int,
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        k_buffer: torch.Tensor,
+        forward_batch: "ForwardBatch",
+    ):
+        """
+        Incrementally update representations during decode phase.
+
+        Called at every layer during forward pass. Algorithm internally decides
+        whether to update based on:
+        - self.states.repr_constructed[req_id]: Whether initial construction done
+        - self.states.last_constructed_page[req_id]: Last constructed page index
+        - Current seq_lens: To detect new tokens/pages
+
+        Algorithm-specific implementations:
+            - ChunkKV: Incrementally compute importance scores for newly generated chunks during decode
+            - Quest: Incrementally compute representations for newly generated pages during decode
+            - PQCache: Assign new tokens to existing centroids (no centroid update during decode)
+            - SnapKV: Optional: periodically re-run voting with sliding observation window (typically static after prefill)
+            - Look-ahead QCache: Periodically regenerate pseudo queries and re-evaluate importance scores to adapt to generation dynamics
+        """
+        pass
+
+    @abstractmethod
+    def retrieve_topk(
+        self,
+        queries: torch.Tensor,
+        layer_id: int,
+        req_pool_indices: torch.Tensor,
+        sparse_mask: torch.Tensor,
+        **kwargs,
+    ) -> tuple:
+        """
+        Retrieve top-k important KV indices for sparse attention.
+
+        Called before attention computation at each layer. Uses current query
+        and pre-computed representations to select the most important subset
+        of KV cache for attention computation.
+
+        Args:
+            queries: [bs, num_heads, head_dim] Current query vectors
+            layer_id: Current layer index
+            req_pool_indices: [bs] Request pool indices
+            sparse_mask: [bs] bool, which requests need sparse attention
+            attn_metadata: Attention metadata (contains seq_lens, etc.)
+            **kwargs: Algorithm-specific arguments
+
+        Returns:
+            selected_indices: [bs, max_selected] Selected page/token indices, padded with -1
+            valid_lengths: [bs] Actual number of selected indices per request
+
+        Note:
+            - Indices are logical positions that will be mapped to physical KV cache by BackendAdaptor
+
+        Algorithm-specific implementations:
+            - ChunkKV: Select top-k chunks based on pre-computed importance scores with layer-wise index reuse
+            - Quest: Compute query-page similarity using current query and stored page representations, select top-k pages
+            - PQCache: Calculate query-centroid similarity, use centroid scores to rank tokens, select top-k tokens
+            - SnapKV: Return union of voted important prefix positions (with clustered neighbors) and observation window tokens
+            - Look-ahead QCache: Return KVs not marked for eviction (eviction based on pseudo query importance evaluation)
+        """
+        pass
+
+
+class BaseSparseAlgorithmImpl(BaseSparseAlgorithm):
+    """
+    Implementation base class for sparse attention algorithms.
+
+    Provides common infrastructure for algorithms that operate at page/chunk granularity
+    (token-wise is simply page_size=1):
+    - Generic construct/update flow with state tracking
+    - TopK retrieval with recent page retention (can be overridden)
+
+    Subclasses need to implement:
+    - _initialize_representation_pools(): Initialize algorithm-specific representation pools
+    - _compute_page_representations(): Compute page scores/representations
+    - _retrieve_page_scores(): Retrieve page scores for TopK selection
+
+    Subclasses can also override any method for specialized behavior
+    """
+
+    def __init__(self, config, device: torch.device, **kwargs):
+        super().__init__(config, device, **kwargs)
+        self.sparsity_ratio = config.sparse_extra_config.get("sparsity_ratio", 0.7)
+        self.num_recent_pages = config.sparse_extra_config.get("num_recent_pages", 4)
+        self.page_size = config.page_size
+
+    def initialize_representation_pool(
+        self,
+        start_layer: int,
+        end_layer: int,
+        token_to_kv_pool,
+        req_to_token_pool,
+        states,
+    ):
+        self.req_to_token_pool = req_to_token_pool
+        self.token_to_kv_pool = token_to_kv_pool
+        self.start_layer = start_layer
+        self.end_layer = end_layer
+        self.states = states
+
+        total_num_tokens = token_to_kv_pool.get_key_buffer(start_layer).shape[0]
+        total_num_pages = (total_num_tokens + self.page_size - 1) // self.page_size
+
+        # Initialize algorithm-specific representation pools
+        self._initialize_representation_pools(start_layer, end_layer, total_num_pages)
+
+    def construct_representations(
+        self,
+        layer_id,
+        req_pool_indices,
+        seq_lens,
+        k_buffer,
+        forward_batch,
+    ) -> torch.Tensor:
+
+        if not forward_batch.forward_mode.is_extend():
+            return
+
+        num_pages = seq_lens // self.page_size
+        valid_mask = (
+            ~self.states.repr_constructed[req_pool_indices]
+            & (seq_lens >= self.states.prompt_lens[req_pool_indices])
+            & (num_pages > 0)
+        )
+
+        if not valid_mask.any():
+            return
+
+        # Compute page representations by subclass
+        self._compute_page_representations(
+            layer_id,
+            req_pool_indices[valid_mask],
+            seq_lens[valid_mask],
+            0,
+            num_pages[valid_mask],
+            k_buffer,
+        )
+
+        # Update tracking states
+        if layer_id == self.end_layer - 1:
+            success_indices = req_pool_indices[valid_mask]
+            self.states.repr_constructed[success_indices] = True
+            self.states.last_constructed_page[success_indices] = num_pages[valid_mask]
+
+    def update_representations(
+        self,
+        layer_id,
+        req_pool_indices,
+        seq_lens,
+        k_buffer,
+        forward_batch,
+    ) -> torch.Tensor:
+        if not forward_batch.forward_mode.is_decode_or_idle():
+            return
+
+        start_page = self.states.last_constructed_page[req_pool_indices]
+        end_page = seq_lens // self.page_size
+        valid_mask = self.states.repr_constructed[req_pool_indices] & (
+            start_page < end_page
+        )
+
+        if not valid_mask.any():
+            return
+
+        # Compute page representations by subclass
+        self._compute_page_representations(
+            layer_id,
+            req_pool_indices[valid_mask],
+            seq_lens[valid_mask],
+            start_page[valid_mask],
+            end_page[valid_mask],
+            k_buffer,
+        )
+
+        # Update tracking states
+        if layer_id == self.end_layer - 1:
+            success_indices = req_pool_indices[valid_mask]
+            self.states.last_constructed_page[success_indices] = end_page[valid_mask]
+
+    def retrieve_topk(
+        self,
+        queries: torch.Tensor,
+        layer_id: int,
+        req_pool_indices: torch.Tensor,
+        sparse_mask: torch.Tensor,
+        **kwargs,
+    ) -> tuple:
+        """
+        Default TopK retrieval: score-based selection + recent pages.
+        Subclasses can override for query-dependent retrieval.
+
+        TODO:
+            1. Using triton kernel to speed up this function
+            2. Support CUDA Graph
+        """
+        bs, device = queries.shape[0], queries.device
+
+        seq_lens_source = kwargs.get("forward_batch", None)
+        if seq_lens_source is None or not hasattr(seq_lens_source, "seq_lens"):
+            raise ValueError(
+                "forward_batch with seq_lens is required for TopK retrieval"
+            )
+        seq_lens = seq_lens_source.seq_lens.to(device)
+
+        req_to_token = self.req_to_token_pool.req_to_token
+        max_req_tokens = req_to_token.shape[1]
+
+        per_request_indices = []
+        per_request_lengths = []
+
+        for i in range(bs):
+            if not sparse_mask[i]:
+                per_request_indices.append(
+                    torch.empty(0, device=device, dtype=torch.int32)
+                )
+                per_request_lengths.append(0)
+                continue
+
+            num_pages = int((seq_lens[i].item() + self.page_size - 1) // self.page_size)
+            if num_pages <= self.num_recent_pages:
+                per_request_indices.append(
+                    torch.empty(0, device=device, dtype=torch.int32)
+                )
+                per_request_lengths.append(0)
+                continue
+
+            page_idx = torch.arange(num_pages, device=device)
+            page_start_token = req_to_token[
+                req_pool_indices[i],
+                (page_idx * self.page_size).clamp(0, max_req_tokens - 1),
+            ]
+            phys_pages = (page_start_token // self.page_size).unsqueeze(0)
+
+            scores = self._retrieve_page_scores(
+                layer_id,
+                phys_pages,
+                req_pool_indices[i : i + 1],
+                queries[i : i + 1],
+            )
+
+            recent_start = max(num_pages - self.num_recent_pages, 0)
+            scores = scores.clone()
+            scores[:, recent_start:] = float("-inf")
+
+            history_pages = max(recent_start, 1)
+            k = max(int(history_pages * self.sparsity_ratio), 1)
+            k = min(k, history_pages)
+            topk_idx = torch.topk(scores, k=k, dim=1, sorted=False)[1].squeeze(0)
+
+            recent_idx = torch.arange(
+                recent_start, recent_start + self.num_recent_pages, device=device
+            )
+            recent_idx = recent_idx[recent_idx < num_pages]
+
+            combined = (
+                torch.cat([topk_idx, recent_idx], dim=0).sort()[0].to(torch.int32)
+            )
+
+            per_request_indices.append(combined)
+            per_request_lengths.append(int(combined.numel()))
+
+        max_len = max(max(per_request_lengths, default=0), 1)
+        out_indices = torch.full((bs, max_len), -1, dtype=torch.int32, device=device)
+        out_lengths = torch.zeros(bs, dtype=torch.int32, device=device)
+
+        for i, selected in enumerate(per_request_indices):
+            length = per_request_lengths[i]
+            if length == 0:
+                continue
+            out_indices[i, :length] = selected
+            out_lengths[i] = length
+
+        return out_indices, out_lengths
+
+    def _initialize_representation_pools(
+        self, start_layer: int, end_layer: int, total_num_pages: int
+    ):
+        """Initialize algorithm-specific representation pools for all layers."""
+        raise NotImplementedError
+
+    def _compute_page_representations(
+        self,
+        layer_id: int,
+        reqs: torch.Tensor,
+        seq_lens: torch.Tensor,
+        start_page,
+        end_page: torch.Tensor,
+        k_buffer: torch.Tensor,
+    ):
+        """Compute and store page representations for given page range."""
+        raise NotImplementedError
+
+    def _retrieve_page_scores(
+        self,
+        layer_id: int,
+        phys_pages: torch.Tensor,
+        req_pool_indices: torch.Tensor,
+        queries: torch.Tensor,
+    ) -> torch.Tensor:
+        """Retrieve page scores for TopK selection."""
+        raise NotImplementedError

python/sglang/srt/mem_cache/sparsity/algorithms/deepseek_dsa.py

@@ -0,0 +1,80 @@
+from typing import Any, Optional
+
+import torch
+
+from sglang.srt.mem_cache.sparsity.algorithms.base_algorithm import (
+    BaseSparseAlgorithmImpl,
+)
+
+
+class DeepSeekDSAAlgorithm(BaseSparseAlgorithmImpl):
+    """
+    Sparse attention algorithm for DeepSeek DSA.
+
+    This algorithm uses DSA's native indexer for TopK retrieval.
+    Overrides all parent methods as DSA has its own specialized flow.
+    """
+
+    def __init__(self, config, device: torch.device, **kwargs):
+        super().__init__(config, device, **kwargs)
+
+    def retrieve_topk(
+        self,
+        queries: torch.Tensor,
+        layer_id: int,
+        req_pool_indices: torch.Tensor,
+        sparse_mask: torch.Tensor,
+        attn_metadata: Optional[Any],
+        **kwargs,
+    ) -> tuple:
+        indexer, forward_batch, x, q_lora, positions = (
+            kwargs.get("indexer"),
+            kwargs.get("forward_batch"),
+            kwargs.get("x"),
+            kwargs.get("q_lora"),
+            kwargs.get("positions"),
+        )
+
+        if any(v is None for v in [indexer, x, q_lora, positions, forward_batch]):
+            raise ValueError("Required: indexer, forward_batch, x, q_lora, positions")
+
+        return (
+            indexer(
+                x=x,
+                q_lora=q_lora,
+                positions=positions,
+                forward_batch=forward_batch,
+                layer_id=layer_id,
+            ),
+            None,
+        )
+
+    def initialize_representation_pool(
+        self,
+        start_layer: int,
+        end_layer: int,
+        token_to_kv_pool,
+        req_to_token_pool,
+        states,
+    ):
+        pass
+
+    def construct_representations(
+        self,
+        layer_id,
+        req_pool_indices,
+        seq_lens,
+        k_buffer,
+        forward_batch,
+    ):
+        pass
+
+    def update_representations(
+        self,
+        layer_id,
+        req_pool_indices,
+        seq_lens,
+        k_buffer,
+        forward_batch,
+    ):
+        pass

python/sglang/srt/mem_cache/sparsity/algorithms/quest_algorithm.py

@@ -0,0 +1,166 @@
+"""
+Quest sparse attention algorithm.
+
+This implementation follows the Quest paper's bounding-box estimation for
+query-aware page selection. For each KV page, it maintains per-dimension
+min/max of keys and uses them to upper-bound attention scores without
+materializing full dot products.
+"""
+
+import logging
+
+import torch
+
+from sglang.srt.mem_cache.sparsity.algorithms.base_algorithm import (
+    BaseSparseAlgorithmImpl,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class QuestAlgorithm(BaseSparseAlgorithmImpl):
+    """Quest page-wise sparse attention using bounding-box criticality."""
+
+    def __init__(self, config, device: torch.device, **kwargs):
+        super().__init__(config, device, **kwargs)
+        self.page_k_min = {}
+        self.page_k_max = {}
+        self.page_valid = {}
+
+    def _initialize_representation_pools(
+        self, start_layer: int, end_layer: int, total_num_pages: int
+    ):
+        key_buf = self.token_to_kv_pool.get_key_buffer(start_layer)
+        head_num, head_dim = key_buf.shape[1], key_buf.shape[2]
+
+        for layer_id in range(start_layer, end_layer):
+            self.page_k_min[layer_id] = torch.zeros(
+                (total_num_pages, head_num, head_dim),
+                dtype=torch.float32,
+                device=self.device,
+            )
+            self.page_k_max[layer_id] = torch.zeros_like(self.page_k_min[layer_id])
+            self.page_valid[layer_id] = torch.zeros(
+                total_num_pages, dtype=torch.bool, device=self.device
+            )
+
+        logger.info(
+            "Initialized Quest page reps: %d pages, %d layers, head_num=%d, head_dim=%d",
+            total_num_pages,
+            end_layer - start_layer,
+            head_num,
+            head_dim,
+        )
+
+    def _compute_page_representations(
+        self,
+        layer_id: int,
+        reqs: torch.Tensor,
+        seq_lens: torch.Tensor,
+        start_page,
+        end_page: torch.Tensor,
+        k_buffer: torch.Tensor,
+    ):
+        if isinstance(start_page, int):
+            start_page = torch.full_like(end_page, start_page)
+
+        device = k_buffer.device
+        req_to_token = self.req_to_token_pool.req_to_token
+        n = reqs.shape[0]
+        max_pages = int((end_page - start_page).max().item())
+        if max_pages <= 0:
+            return
+
+        pg_off = torch.arange(max_pages, device=device).unsqueeze(0)
+        pg_id = start_page.unsqueeze(1) + pg_off
+        pg_mask = pg_id < end_page.unsqueeze(1)
+
+        tok_start = pg_id * self.page_size
+        tok_off = torch.arange(self.page_size, device=device).view(1, 1, -1)
+        tok_pos = tok_start.unsqueeze(2) + tok_off
+        tok_mask = (
+            tok_pos
+            < (tok_start + self.page_size).clamp(max=seq_lens.unsqueeze(1)).unsqueeze(2)
+        ) & pg_mask.unsqueeze(2)
+
+        phys_tok = req_to_token[
+            reqs.view(n, 1, 1).expand(n, max_pages, self.page_size),
+            tok_pos.clamp(0, req_to_token.shape[1] - 1),
+        ].clamp(0, k_buffer.shape[0] - 1)
+
+        keys = k_buffer[phys_tok].to(torch.float32)
+        mask = tok_mask.unsqueeze(-1).unsqueeze(-1)
+
+        page_min = torch.where(mask, keys, torch.full_like(keys, float("inf"))).amin(
+            dim=2
+        )
+        page_max = torch.where(mask, keys, torch.full_like(keys, float("-inf"))).amax(
+            dim=2
+        )
+
+        phys_pg = (
+            req_to_token[
+                reqs.unsqueeze(1).expand(n, max_pages),
+                tok_start.clamp(0, req_to_token.shape[1] - 1),
+            ]
+            // self.page_size
+        )
+
+        idx = pg_mask.nonzero(as_tuple=False)
+        if idx.numel() == 0:
+            return
+
+        target_pages = phys_pg[idx[:, 0], idx[:, 1]].clamp(
+            0, self.page_k_min[layer_id].shape[0] - 1
+        )
+        self.page_k_min[layer_id][target_pages] = page_min[idx[:, 0], idx[:, 1]]
+        self.page_k_max[layer_id][target_pages] = page_max[idx[:, 0], idx[:, 1]]
+        self.page_valid[layer_id][target_pages] = True
+
+    def _retrieve_page_scores(
+        self,
+        layer_id: int,
+        phys_pages: torch.Tensor,
+        req_pool_indices: torch.Tensor,
+        queries: torch.Tensor,
+    ) -> torch.Tensor:
+        # Clamp pages to valid storage range
+        phys_pages_clamped = phys_pages.clamp(0, self.page_k_min[layer_id].shape[0] - 1)
+
+        k_min = self.page_k_min[layer_id][phys_pages_clamped]
+        k_max = self.page_k_max[layer_id][phys_pages_clamped]
+        valid_mask = self.page_valid[layer_id][phys_pages_clamped]
+        # Align query shape to KV heads.
+        head_dim = k_min.shape[-1]
+        if queries.dim() == 2:
+            bs, hidden = queries.shape
+            if hidden % head_dim != 0:
+                raise ValueError(
+                    f"Quest query hidden size {hidden} not divisible by head_dim {head_dim}"
+                )
+            q_heads = hidden // head_dim
+            q = queries.view(bs, q_heads, head_dim)
+        elif queries.dim() == 3:
+            q = queries
+        else:
+            raise ValueError(f"Unsupported query shape for Quest: {queries.shape}")
+
+        kv_heads = k_min.shape[-2]
+        q_heads = q.shape[1]
+        if q_heads != kv_heads:
+            if q_heads % kv_heads != 0:
+                raise ValueError(
+                    f"Query heads {q_heads} not divisible by KV heads {kv_heads}"
+                )
+            group = q_heads // kv_heads
+            # Average grouped query heads to align with KV heads (approximation for MQA/GQA).
+            q = q.view(q.shape[0], kv_heads, group, head_dim).mean(dim=2)
+
+        q = q.to(k_min.dtype).unsqueeze(1)  # [bs, 1, kv_heads, head_dim]
+
+        criticality = torch.where(q >= 0, q * k_max, q * k_min).sum(dim=(2, 3))
+        criticality = torch.where(
+            valid_mask, criticality, torch.full_like(criticality, float("-inf"))
+        )
+
+        return criticality

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)

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

@@ -0,0 +1,11 @@
+from sglang.srt.mem_cache.sparsity.core.sparse_coordinator import (
+    RequestTrackers,
+    SparseConfig,
+    SparseCoordinator,
+)
+
+__all__ = [
+    "RequestTrackers",
+    "SparseConfig",
+    "SparseCoordinator",
+]

python/sglang/srt/mem_cache/sparsity/core/sparse_coordinator.py

@@ -0,0 +1,276 @@
+import logging
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Optional
+
+import torch
+
+from sglang.srt.mem_cache.memory_pool import KVCache, ReqToTokenPool
+from sglang.srt.mem_cache.sparsity.algorithms.base_algorithm import BaseSparseAlgorithm
+from sglang.srt.mem_cache.sparsity.backend.backend_adaptor import BackendAdaptor
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.radix_attention import RadixAttention
+    from sglang.srt.managers.schedule_batch import Req
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+
+logger = logging.getLogger(__name__)
+
+
+class RequestTrackers:
+    """State tracker for sparse attention requests."""
+
+    def __init__(
+        self,
+        max_pool_size: int,
+        device: torch.device,
+        num_layers: int,
+        min_sparse_prompt_len: int,
+        max_context_len: int,
+    ):
+        self.device = device
+        self.num_layers = num_layers
+
+        self.repr_constructed = torch.zeros(
+            max_pool_size, dtype=torch.bool, device=device
+        )
+        self.prompt_lens = torch.zeros(max_pool_size, dtype=torch.int64, device=device)
+        self.last_constructed_page = torch.zeros(
+            max_pool_size, dtype=torch.int64, device=device
+        )
+
+        # TODO: Add more trackers for hierarchical KVCache management
+
+    def register(self, idx: int, prompt_len: int) -> None:
+        self.repr_constructed[idx] = False
+        self.prompt_lens[idx] = prompt_len
+        self.last_constructed_page[idx] = 0
+
+    def clear(self, idx: int) -> None:
+        self.repr_constructed[idx] = False
+        self.prompt_lens[idx] = 0
+        self.last_constructed_page[idx] = 0
+
+
+@dataclass
+class SparseConfig:
+    """Configuration for sparse attention."""
+
+    top_k: int = 2048
+    device_buffer_size: int = 4096
+    host_to_device_ratio: int = 2
+    swap_in_block_size: int = 960
+    algorithm: Optional[str] = None
+    backend: Optional[str] = None
+    page_size: Optional[int] = None
+    min_sparse_prompt_len: Optional[int] = None
+    sparse_extra_config: dict = field(
+        default_factory=dict
+    )  # Algorithm-specific config, parsed by each algorithm
+
+
+class SparseCoordinator:
+    """
+    Coordinator for sparse attention with retrievable KV cache compression.
+
+    This coordinator framework is designed for decode-phase retrievable algorithms
+    (e.g., Quest, PQCache, SnapKV) that dynamically select important KV cache entries
+    based on current queries. It manages the lifecycle of sparse attention including
+    representation construction, sparse retrieval, and token offloading.
+
+    Request Lifecycle and API Calls:
+        1. Request Start:
+           - on_request_begin(req) -> Register request and initialize state
+
+        2. Prefill Phase:
+           - attention_end(...)    -> Construct representations
+
+        3. Decode Phase:
+           - forward_begin(batch)  -> Wait for pending KVCache offloading
+           - attention_begin(...)  -> Identify important KV, load offloaded KVCache, adapt attention metadata
+           - attention_end(...)    -> Construct/update representations
+           - forward_end(batch)    -> Trigger KVCache offloading
+
+        4. Request End:
+           - on_request_end(req) -> Clean up state and resources
+    """
+
+    def __init__(
+        self,
+        config: SparseConfig,
+        algorithm: BaseSparseAlgorithm,
+        backend_adaptor: Optional[BackendAdaptor],
+        req_to_token_pool: ReqToTokenPool,
+        token_to_kv_pool: KVCache,
+        start_layer: int,
+        end_layer: int,
+        device: torch.device,
+    ):
+        self.config = config
+        self.algorithm = algorithm
+        self.backend_adaptor = backend_adaptor
+        self.req_to_token_pool = req_to_token_pool
+        self.token_to_kv_pool = token_to_kv_pool
+        self.start_layer = start_layer
+        self.end_layer = end_layer
+        self.device = device
+        self.page_size = config.page_size
+
+        self.states = RequestTrackers(
+            req_to_token_pool.req_to_token.shape[0],
+            device,
+            end_layer - start_layer + 1,
+            self.config.min_sparse_prompt_len,
+            self.req_to_token_pool.max_context_len,
+        )
+
+        # Initialize algorithm representation pool and context
+        self.algorithm.initialize_representation_pool(
+            start_layer,
+            end_layer,
+            self.token_to_kv_pool,
+            self.req_to_token_pool,
+            self.states,
+        )
+
+        logger.info(
+            f"SparseCoordinator initialized with sparse algorithm={type(algorithm).__name__}"
+        )
+
+    def on_request_begin(self, req: "Req") -> None:
+        """
+        Handle request begin event. Called when a new request is created.
+
+        Registers the request in the state tracker to enable sparse attention processing.
+        """
+        if req.req_pool_idx is not None:
+            self.states.register(req.req_pool_idx, len(req.origin_input_ids))
+
+    def on_request_end(self, req: "Req") -> None:
+        """
+        Handle request end event. Called when a request is completed or aborted.
+        Cleans up request-specific state and releases resources.
+        """
+        if req.req_pool_idx is None:
+            return
+
+        self.states.clear(req.req_pool_idx)
+
+        # TODO: Implement request end handling
+        # - Release host indices if any were allocated for offloading
+
+    def forward_begin(self, forward_batch: "ForwardBatch") -> None:
+        """
+        Handle forward pass begin event. Called before each forward pass starts.
+
+        Wait for pending KVCache offloading operations to complete before forward pass.
+        Ensures memory consistency for subsequent sparse attention operations.
+        """
+        # TODO: Implement forward begin handling
+        # - Check if there are pending offloading operations
+        pass
+
+    def forward_end(self, forward_batch: "ForwardBatch") -> None:
+        """
+        Handle forward pass end event. Called after each forward pass completes.
+
+        Trigger async KVCache offloading operations.
+        """
+        # TODO: Implement forward end handling
+        # - Identify tokens to offload
+        # - Trigger async offloading operations
+        pass
+
+    def attention_begin(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        layer: "RadixAttention",
+        forward_batch: "ForwardBatch",
+        attn_metadata: Optional[Any],
+        **kwargs,
+    ) -> Optional[Any]:
+        """
+        Handle attention begin event. Called before each attention pass starts.
+
+        Identify important KV entries via sparse algorithm, load offloaded KVCache if needed,
+        and adapt attention metadata for the attention backend.
+        """
+        if layer.layer_id == self.start_layer:
+            self.backend_adaptor.save_original_metadata(attn_metadata)
+
+        return self._handle_sparse_retrieve(
+            query, layer, forward_batch, attn_metadata, **kwargs
+        )
+
+    def attention_end(
+        self,
+        output: torch.Tensor,
+        layer: "RadixAttention",
+        forward_batch: "ForwardBatch",
+    ) -> None:
+        """
+        Handle attention end event. Called after each attention pass completes.
+
+        Maybe construct and update sparse representations.
+        """
+        layer_id = layer.layer_id
+
+        # Maybe construct representations
+        self.algorithm.construct_representations(
+            layer_id=layer_id,
+            req_pool_indices=forward_batch.req_pool_indices,
+            seq_lens=forward_batch.seq_lens,
+            k_buffer=self.token_to_kv_pool.get_key_buffer(layer_id),
+            forward_batch=forward_batch,
+        )
+
+        # Maybe update representations
+        self.algorithm.update_representations(
+            layer_id=layer_id,
+            req_pool_indices=forward_batch.req_pool_indices,
+            seq_lens=forward_batch.seq_lens,
+            k_buffer=self.token_to_kv_pool.get_key_buffer(layer_id),
+            forward_batch=forward_batch,
+        )
+
+    def _handle_sparse_retrieve(
+        self,
+        query: torch.Tensor,
+        layer: "RadixAttention",
+        forward_batch: "ForwardBatch",
+        attn_metadata: Optional[Any],
+        **kwargs,
+    ) -> Optional[torch.Tensor]:
+        req_pool_indices = forward_batch.req_pool_indices
+        # Compute Topk
+        sparse_mask = self._compute_sparse_mask(req_pool_indices)
+        selected_indices, valid_lengths = self.algorithm.retrieve_topk(
+            queries=query,
+            layer_id=layer.layer_id,
+            req_pool_indices=req_pool_indices,
+            sparse_mask=sparse_mask,
+            forward_batch=forward_batch,
+            attn_metadata=attn_metadata,
+            **kwargs,
+        )
+
+        # Adapt Attention Metadata
+        return self.backend_adaptor.adapt_for_attn_metadata(
+            selected_indices=selected_indices,
+            valid_lengths=valid_lengths,
+            sparse_mask=sparse_mask,
+            current_metadata=attn_metadata,
+            forward_batch=forward_batch,
+            req_to_token=self.req_to_token_pool.req_to_token,
+            page_size=self.page_size,
+            layer_id=layer.layer_id,
+        )
+
+    def _compute_sparse_mask(self, req_pool_indices):
+        mask = (
+            self.states.prompt_lens[req_pool_indices]
+            >= self.config.min_sparse_prompt_len
+        )
+
+        return mask

python/sglang/srt/mem_cache/sparsity/factory.py

@@ -0,0 +1,154 @@
+import json
+import logging
+from typing import Optional
+
+import torch
+
+from sglang.srt.mem_cache.sparsity.algorithms.base_algorithm import BaseSparseAlgorithm
+from sglang.srt.mem_cache.sparsity.algorithms.deepseek_dsa import DeepSeekDSAAlgorithm
+from sglang.srt.mem_cache.sparsity.algorithms.quest_algorithm import QuestAlgorithm
+from sglang.srt.mem_cache.sparsity.backend.backend_adaptor import (
+    DSABackendAdaptor,
+    FlashAttentionAdaptor,
+)
+from sglang.srt.mem_cache.sparsity.core.sparse_coordinator import (
+    SparseConfig,
+    SparseCoordinator,
+)
+
+logger = logging.getLogger(__name__)
+
+_global_sparse_coordinator: Optional[SparseCoordinator] = None
+
+_ALGORITHM_REGISTRY = {
+    "quest": lambda config, device, **kw: QuestAlgorithm(config, device, **kw),
+    "deepseek_dsa": lambda config, device, **kw: DeepSeekDSAAlgorithm(
+        config, device, **kw
+    ),
+}
+
+
+def _create_sparse_algorithm(
+    config: SparseConfig,
+    device: torch.device,
+    **kwargs,
+) -> BaseSparseAlgorithm:
+    algorithm_name = config.algorithm.lower()
+    factory = _ALGORITHM_REGISTRY.get(algorithm_name)
+
+    if factory is None:
+        raise ValueError(f"Unknown sparse algorithm: {algorithm_name}")
+
+    return factory(config, device, **kwargs)
+
+
+def _create_backend_adaptor(
+    backend: str,
+    device: torch.device,
+    sparse_algorithm: BaseSparseAlgorithm,
+    req_to_token_pool,
+):
+    """Create backend adaptor."""
+    if isinstance(sparse_algorithm, DeepSeekDSAAlgorithm):
+        return DSABackendAdaptor(device, req_to_token_pool)
+
+    if backend in ["fa3", "flashattention"]:
+        return FlashAttentionAdaptor(device)
+
+    raise ValueError(f"Unknown attention backend: {backend}")
+
+
+def _parse_sparse_config(server_args) -> SparseConfig:
+    """Parse hierarchical sparse config from JSON string.
+
+    Required fields with defaults: top_k (2048), device_buffer_size (2*top_k),
+    host_to_device_ratio (2), swap_in_block_size (960).
+    Optional fields (default None): algorithm, backend, min_sparse_prompt_len,
+    page_size. All remaining fields go to sparse_extra_config.
+    """
+    extra_config_str = server_args.hisparse_config
+    if extra_config_str is not None:
+        try:
+            extra_config = json.loads(extra_config_str)
+        except json.JSONDecodeError as e:
+            raise ValueError(f"Failed to parse hisparse_config: {e}") from e
+    else:
+        extra_config = {}
+
+    top_k = extra_config.pop("top_k", 2048)
+    device_buffer_size = extra_config.pop("device_buffer_size", 2 * top_k)
+    host_to_device_ratio = extra_config.pop("host_to_device_ratio", 2)
+    swap_in_block_size = extra_config.pop("swap_in_block_size", 960)
+
+    if device_buffer_size < top_k:
+        raise ValueError(
+            f"device_buffer_size ({device_buffer_size}) must be no smaller than top_k ({top_k})"
+        )
+    if not isinstance(swap_in_block_size, int) or isinstance(swap_in_block_size, bool):
+        raise ValueError(
+            f"swap_in_block_size must be an integer, got {swap_in_block_size!r}"
+        )
+    if swap_in_block_size <= 0 or swap_in_block_size > 1024:
+        raise ValueError(
+            f"swap_in_block_size ({swap_in_block_size}) must be in the range [1, 1024]"
+        )
+
+    algorithm = extra_config.pop("algorithm", None)
+    backend = extra_config.pop("backend", None)
+    min_sparse_prompt_len = extra_config.pop("min_sparse_prompt_len", None)
+    page_size = extra_config.pop("page_size", None)
+
+    return SparseConfig(
+        top_k=top_k,
+        device_buffer_size=device_buffer_size,
+        host_to_device_ratio=host_to_device_ratio,
+        swap_in_block_size=swap_in_block_size,
+        algorithm=algorithm,
+        backend=backend,
+        page_size=page_size,
+        min_sparse_prompt_len=min_sparse_prompt_len,
+        sparse_extra_config=extra_config,
+    )
+
+
+def parse_hisparse_config(server_args) -> SparseConfig:
+    """Parse hisparse config from server_args, returning defaults if no config provided."""
+    return _parse_sparse_config(server_args)
+
+
+def create_sparse_coordinator(
+    device: torch.device,
+    req_to_token_pool,
+    token_to_kv_pool,
+    start_layer: int,
+    end_layer: int,
+    server_args,
+    **kwargs,
+) -> SparseCoordinator:
+    config = _parse_sparse_config(server_args)
+    algorithm = _create_sparse_algorithm(config, device, **kwargs)
+    backend_adaptor = _create_backend_adaptor(
+        config.backend, device, algorithm, req_to_token_pool
+    )
+
+    coordinator = SparseCoordinator(
+        config=config,
+        algorithm=algorithm,
+        backend_adaptor=backend_adaptor,
+        req_to_token_pool=req_to_token_pool,
+        token_to_kv_pool=token_to_kv_pool,
+        start_layer=start_layer,
+        end_layer=end_layer,
+        device=device,
+    )
+    register_sparse_coordinator(coordinator)
+    return coordinator
+
+
+def register_sparse_coordinator(coordinator: SparseCoordinator) -> None:
+    global _global_sparse_coordinator
+    _global_sparse_coordinator = coordinator
+
+
+def get_sparse_coordinator() -> Optional[SparseCoordinator]:
+    return _global_sparse_coordinator