[3/N][Sparse With Hicache]: Init sparse coordinator (#16086)

Co-authored-by: 晟海 <huangtingwei.htw@antgroup.com>
Co-authored-by: Zhiqiang Xie <xiezhq@stanford.edu>

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

@@ -0,0 +1,27 @@
+from sglang.srt.mem_cache.sparsity.algorithms import (
+    BaseSparseAlgorithm,
+    BaseSparseAlgorithmImpl,
+    DeepSeekNSAAlgorithm,
+    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,
+    register_sparse_coordinator,
+)
+
+__all__ = [
+    "BaseSparseAlgorithm",
+    "BaseSparseAlgorithmImpl",
+    "QuestAlgorithm",
+    "DeepSeekNSAAlgorithm",
+    "BackendAdaptor",
+    "FlashAttentionAdaptor",
+    "SparseConfig",
+    "SparseCoordinator",
+    "create_sparse_coordinator",
+    "get_sparse_coordinator",
+    "register_sparse_coordinator",
+]

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

@@ -162,9 +162,9 @@ class BaseSparseAlgorithmImpl(BaseSparseAlgorithm):
 
     def __init__(self, config, device: torch.device, **kwargs):
         super().__init__(config, device, **kwargs)
-        self.compression_ratio = getattr(config, "compression_ratio", 0.3)
-        self.page_size = getattr(config, "page_size", 64)
-        self.num_recent_pages = getattr(config, "num_recent_pages", 4)
+        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,
@@ -325,7 +325,7 @@ class BaseSparseAlgorithmImpl(BaseSparseAlgorithm):
             scores[:, recent_start:] = float("-inf")
 
             history_pages = max(recent_start, 1)
-            k = max(int(history_pages * (1 - self.compression_ratio)), 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)
 

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

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

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 NSABackendAdaptor(BackendAdaptor):
+    """Adaptor for NSA (Native 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 NSA backend.
+        """
+        # TODO: Implement NSA 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,272 @@
+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."""
+
+    backend: str
+    algorithm: str
+    page_size: int = 64
+    min_sparse_prompt_len: int = 2048
+    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,126 @@
+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_nsa import DeepSeekNSAAlgorithm
+from sglang.srt.mem_cache.sparsity.algorithms.quest_algorithm import QuestAlgorithm
+from sglang.srt.mem_cache.sparsity.backend.backend_adaptor import (
+    FlashAttentionAdaptor,
+    NSABackendAdaptor,
+)
+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_nsa": lambda config, device, **kw: DeepSeekNSAAlgorithm(
+        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, DeepSeekNSAAlgorithm):
+        return NSABackendAdaptor(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"""
+    # Parse extra config if provided
+    extra_config_str = server_args.hierarchical_sparse_attention_extra_config
+    if extra_config_str is not None:
+        try:
+            extra_config = json.loads(extra_config_str)
+
+            # Extract algorithm and backend
+            algorithm = extra_config.pop("algorithm", "quest")
+            backend = extra_config.pop("backend", "flashattention")
+            min_sparse_prompt_len = extra_config.pop("min_sparse_prompt_len", 2048)
+
+            # Everything else goes to algorithm_extra_config
+            sparse_extra_config = extra_config
+        except json.JSONDecodeError as e:
+            logger.warning(
+                f"Failed to parse hierarchical_sparse_attention_extra_config: {e}"
+            )
+
+    config = SparseConfig(
+        algorithm=algorithm,
+        backend=backend,
+        page_size=server_args.page_size,
+        min_sparse_prompt_len=min_sparse_prompt_len,
+        sparse_extra_config=sparse_extra_config,
+    )
+    return config
+
+
+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

python/sglang/srt/server_args.py

@@ -483,6 +483,10 @@ class ServerArgs:
     hicache_storage_backend: Optional[str] = None
     hicache_storage_prefetch_policy: str = "best_effort"
     hicache_storage_backend_extra_config: Optional[str] = None
+
+    # Hierarchical sparse attention
+    hierarchical_sparse_attention_extra_config: Optional[str] = None
+
     # LMCache
     enable_lmcache: bool = False
 
@@ -3816,6 +3820,18 @@ class ServerArgs:
             default=ServerArgs.hicache_storage_backend_extra_config,
             help="A dictionary in JSON string format containing extra configuration for the storage backend.",
         )
+
+        # Hierarchical sparse attention
+        parser.add_argument(
+            "--hierarchical-sparse-attention-extra-config",
+            type=str,
+            default=ServerArgs.hierarchical_sparse_attention_extra_config,
+            help="A dictionary in JSON string format for hierarchical sparse attention configuration. "
+            "Required fields: algorithm (str), backend (str). "
+            "All other fields are algorithm-specific and passed to the algorithm constructor. "
+            'Example: \'{"algorithm": "quest", "backend": "flashattention", "sparsity_ratio": 0.7, "min_sparse_prompt_len": 2048}\'',
+        )
+
         # LMCache
         parser.add_argument(
             "--enable-lmcache",