[1/N][Sparse With Hicache]: Add Sparse Interface

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

@@ -0,0 +1,17 @@
+from sglang.srt.mem_cache.sparsity.algorithms.base_algorithm import (
+    BaseSparseAlgorithm,
+    SparseMode,
+)
+from sglang.srt.mem_cache.sparsity.algorithms.deepseek_nsa import DeepSeekNSAAlgorithm
+from sglang.srt.mem_cache.sparsity.algorithms.page_wise_algorithm import (
+    BasePageWiseAlgorithm,
+    KnormPageAlgorithm,
+)
+
+__all__ = [
+    "BaseSparseAlgorithm",
+    "SparseMode",
+    "BasePageWiseAlgorithm",
+    "KnormPageAlgorithm",
+    "DeepSeekNSAAlgorithm",
+]

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

@@ -0,0 +1,175 @@
+from abc import ABC, abstractmethod
+from enum import Enum
+from typing import TYPE_CHECKING, Any, Optional
+
+import torch
+
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+
+
+class SparseMode(Enum):
+    """Sparse attention granularity mode."""
+
+    PAGE_WISE = "page_wise"
+    TOKEN_WISE = "token_wise"
+    DEEPSEEK_TOKEN_WISE = "deepseek_token_wise"
+
+
+class BaseSparseAlgorithm(ABC):
+    """
+    Abstract base class for sparse attention algorithms.
+
+    This class provides a unified interface for implementing various retrievable KVCache
+    compression algorithms, supporting both page-wise and token-wise sparsity.
+
+    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]
+        """
+        self.req_to_token_pool = req_to_token_pool
+        self.states = states
+
+    @abstractmethod
+    def get_sparse_mode(self) -> SparseMode:
+        """
+        Return the sparsity granularity mode.
+
+        Returns:
+            SparseMode.PAGE_WISE: Selection operates on page/chunk level
+            SparseMode.TOKEN_WISE: Selection operates on individual token level
+
+        Algorithm-specific modes:
+            - ChunkKV: PAGE_WISE (selects important semantic chunks while preserving linguistic structures)
+            - Quest: PAGE_WISE (selects important pages/blocks)
+            - PQCache: TOKEN_WISE (selects important tokens via centroid similarity)
+            - SnapKV: TOKEN_WISE (retention-based: keeps voted important prefix tokens + observation window)
+            - Look-ahead QCache: TOKEN_WISE (eviction-based: removes tokens with low pseudo query importance)
+        """
+        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 based on self.states.repr_constructed.
+        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_extracted_token[req_id]: Last processed position
+        - 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,
+        attn_metadata: Optional[Any],
+        **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
+
+        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)
+
+        Note:
+            - For PAGE_WISE mode: Returns page indices
+            - For TOKEN_WISE mode: Returns token indices
+            - Indices are logical positions that will be mapped to physical KV cache by BackendAdaptor
+        """
+        pass

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

@@ -0,0 +1,74 @@
+import logging
+from typing import TYPE_CHECKING, Any, Optional
+
+import nvtx
+import torch
+
+from sglang.srt.mem_cache.sparsity.algorithms.base_algorithm import (
+    BaseSparseAlgorithm,
+    SparseMode,
+)
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+
+logger = logging.getLogger(__name__)
+
+
+class DeepSeekNSAAlgorithm(BaseSparseAlgorithm):
+    """Sparse attention algorithm for DeepSeek NSA."""
+
+    def __init__(self, config, device: torch.device, **kwargs):
+        super().__init__(config, device, **kwargs)
+        self.index_topk = getattr(config, "index_topk", 2048)
+        logger.info(f"DeepSeekNSAAlgorithm initialized: index_topk={self.index_topk}")
+
+    def get_sparse_mode(self) -> SparseMode:
+        return SparseMode.DEEPSEEK_TOKEN_WISE
+
+    @nvtx.annotate("DeepSeekNSAAlgorithm.retrieve_topk", color="green")
+    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: Optional["Indexer"] = kwargs.get("indexer")
+        forward_batch: Optional["ForwardBatch"] = kwargs.get("forward_batch")
+        x, q_lora, positions = (
+            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, x, q_lora, positions, forward_batch")
+
+        try:
+            # Using the nsa's original indexer to get the topk indices.
+            topk_indices = indexer(
+                x=x,
+                q_lora=q_lora,
+                positions=positions,
+                forward_batch=forward_batch,
+                layer_id=layer_id,
+            )
+
+            if topk_indices is None:
+                return self._empty_result(queries.shape[0], queries.device)
+
+            return topk_indices, None
+        except Exception as e:
+            logger.error(f"Layer {layer_id} NSA indexer failed: {e}", exc_info=True)
+            return self._empty_result(queries.shape[0], queries.device)
+
+    def _empty_result(self, batch_size: int, device: torch.device) -> tuple:
+        selected_indices = torch.full(
+            (batch_size, self.index_topk), -1, dtype=torch.int32, device=device
+        )
+        valid_lengths = torch.zeros(batch_size, dtype=torch.int32, device=device)
+        return selected_indices, valid_lengths