Remove deprecated double sparsity feature (#23009)

docs/advanced_features/server_arguments.md

@@ -382,16 +382,6 @@ Please consult the documentation below and [server_args.py](https://github.com/s
 | `--dllm-algorithm` | The diffusion LLM algorithm, such as LowConfidence. | `None` | Type: str |
 | `--dllm-algorithm-config` | The diffusion LLM algorithm configurations. Must be a YAML file. | `None` | Type: str |
 
-## Double Sparsity
-| Argument | Description | Defaults | Options |
-| --- | --- | --- | --- |
-| `--enable-double-sparsity` | Enable double sparsity attention | `False` | bool flag (set to enable) |
-| `--ds-channel-config-path` | The path of the double sparsity channel config | `None` | Type: str |
-| `--ds-heavy-channel-num` | The number of heavy channels in double sparsity attention | `32` | Type: int |
-| `--ds-heavy-token-num` | The number of heavy tokens in double sparsity attention | `256` | Type: int |
-| `--ds-heavy-channel-type` | The type of heavy channels in double sparsity attention | `qk` | Type: str |
-| `--ds-sparse-decode-threshold` | The minimum decode sequence length required before the double-sparsity backend switches from the dense fallback to the sparse decode kernel. | `4096` | Type: int |
-
 ## Offloading
 | Argument | Description | Defaults | Options |
 | --- | --- | --- | --- |

docs/platforms/ascend/ascend_npu_support_features.md

@@ -480,10 +480,4 @@ The following parameters have some functional deficiencies on community
 
 | Argument                              | Defaults | Options                        |
 |---------------------------------------|----------|--------------------------------|
-| `--enable-double-sparsity`            | `False`  | bool flag<br/> (set to enable) |
-| `--ds-channel-config-path`            | `None`   | Type: str                      |
-| `--ds-heavy-channel-num`              | `32`     | Type: int                      |
-| `--ds-heavy-token-num`                | `256`    | Type: int                      |
-| `--ds-heavy-channel-type`             | `qk`     | Type: str                      |
-| `--ds-sparse-decode-`<br/>`threshold` | `4096`   | Type: int                      |
 | `--tool-server`                       | `None`   | Type: str                      |

python/sglang/srt/layers/attention/attention_registry.py

@@ -98,16 +98,9 @@ def create_triton_backend(runner):
         "Cross attention is not supported in the triton attention backend. "
         "Please use `--attention-backend flashinfer`."
     )
-    if runner.server_args.enable_double_sparsity:
+    from sglang.srt.layers.attention.triton_backend import TritonAttnBackend
-        from sglang.srt.layers.attention.double_sparsity_backend import (
-            DoubleSparseAttnBackend,
-        )
 
-        return DoubleSparseAttnBackend(runner)
+    return TritonAttnBackend(runner)
-    else:
-        from sglang.srt.layers.attention.triton_backend import TritonAttnBackend
-
-        return TritonAttnBackend(runner)
 
 
 @register_attention_backend("torch_native")

python/sglang/srt/layers/attention/double_sparsity_backend.py

@@ -1,257 +0,0 @@
-from __future__ import annotations
-
-from typing import TYPE_CHECKING
-
-import torch
-
-from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
-from sglang.srt.model_executor.forward_batch_info import ForwardBatch
-from sglang.srt.server_args import get_global_server_args
-
-if TYPE_CHECKING:
-    from sglang.srt.layers.radix_attention import RadixAttention
-    from sglang.srt.model_executor.model_runner import ModelRunner
-
-
-class DoubleSparseAttnBackend(AttentionBackend):
-    def __init__(self, model_runner: ModelRunner):
-        # Lazy import to avoid the initialization of cuda context
-        from sglang.srt.layers.attention.triton_ops.double_sparsity_attention import (
-            extend_attention_fwd,
-            flash_decode_attention_fwd,
-            flash_decode_sparse_attention_fwd,
-        )
-
-        super().__init__()
-
-        self.decode_attention_fwd = flash_decode_attention_fwd
-        self.decode_sparse_attention_fwd = flash_decode_sparse_attention_fwd
-        self.extend_attention_fwd = extend_attention_fwd
-        self.num_head = model_runner.model_config.num_attention_heads
-        self.head_dim = model_runner.model_config.hidden_size // self.num_head
-        self.heavy_token_num = model_runner.server_args.ds_heavy_token_num
-
-        self.sorted_channels = model_runner.sorted_channels
-        self.sparse_decode_threshold = (
-            model_runner.server_args.ds_sparse_decode_threshold
-        )
-        self.att_out_approx: torch.Tensor = None
-        self.mid_out: torch.Tensor = None
-        self.mid_o_logexpsum: torch.Tensor = None
-
-        # TODO: Change the hard-coded block_seq_num
-        self.BLOCK_SEQ = 128
-
-        if get_global_server_args().triton_attention_reduce_in_fp32:
-            self.reduce_dtype = torch.float32
-        else:
-            self.reduce_dtype = torch.float16
-
-        self.forward_metadata = None
-
-    def init_forward_metadata(self, forward_batch: ForwardBatch):
-        """Init auxiliary variables for triton attention backend."""
-
-        if forward_batch.forward_mode.is_decode():
-            start_loc = torch.zeros_like(forward_batch.seq_lens, dtype=torch.int32)
-            start_loc[1:] = torch.cumsum(forward_batch.seq_lens[:-1], dim=0)
-
-            total_num_tokens = torch.sum(forward_batch.seq_lens).item()
-            attn_logits = torch.empty(
-                (self.num_head, total_num_tokens),
-                dtype=self.reduce_dtype,
-                device="cuda",
-            )
-
-            max_seq_len = torch.max(forward_batch.seq_lens).item()
-            min_seq_len = torch.min(forward_batch.seq_lens).item()
-            max_extend_len = None
-            # NOTE: Align sequence order with req_to_token order
-            ds_req_to_token = forward_batch.req_to_token_pool.req_to_token[
-                forward_batch.req_pool_indices
-            ]
-
-            bsz = forward_batch.seq_lens.shape[0]
-
-            att_out_approx = torch.empty(
-                [self.num_head, bsz, max_seq_len],
-                dtype=self.reduce_dtype,
-                device="cuda",
-            )
-
-            block_seq_num = (
-                self.heavy_token_num + self.BLOCK_SEQ - 1
-            ) // self.BLOCK_SEQ
-
-            mid_out = torch.empty(
-                [bsz, self.num_head, block_seq_num, self.head_dim],
-                dtype=torch.float32,
-                device="cuda",
-            )
-            mid_o_logexpsum = torch.empty(
-                [bsz, self.num_head, block_seq_num], dtype=torch.float32, device="cuda"
-            )
-            self.att_out_approx = att_out_approx
-            self.mid_out = mid_out
-            self.mid_o_logexpsum = mid_o_logexpsum
-
-        else:
-            start_loc = attn_logits = max_seq_len = min_seq_len = None
-            prefix_lens = forward_batch.extend_prefix_lens
-            max_extend_len = torch.max(forward_batch.seq_lens - prefix_lens).item()
-            ds_req_to_token = None
-
-        self.forward_metadata = (
-            start_loc,
-            attn_logits,
-            max_seq_len,
-            min_seq_len,
-            max_extend_len,
-            ds_req_to_token,
-        )
-
-    def forward_extend(
-        self,
-        q,
-        k,
-        v,
-        layer: RadixAttention,
-        forward_batch: ForwardBatch,
-        save_kv_cache=True,
-    ):
-        # TODO: reuse the buffer across layers
-        if layer.qk_head_dim != layer.v_head_dim:
-            o = q.new_empty((q.shape[0], layer.tp_q_head_num * layer.v_head_dim))
-        else:
-            o = torch.empty_like(q)
-
-        k_label = torch.gather(
-            k,
-            2,
-            self.sorted_channels[layer.layer_id]
-            .unsqueeze(0)
-            .expand(k.shape[0], -1, -1),
-        )
-
-        if save_kv_cache:
-            forward_batch.token_to_kv_pool.set_kv_buffer(
-                layer, forward_batch.out_cache_loc, k, v, k_label
-            )
-
-        (
-            start_loc,
-            attn_logits,
-            max_seq_len,
-            min_seq_len,
-            max_extend_len,
-            ds_req_to_token,
-        ) = self.forward_metadata
-        self.extend_attention_fwd(
-            q.view(-1, layer.tp_q_head_num, layer.qk_head_dim),
-            k.contiguous(),
-            v.contiguous(),
-            o.view(-1, layer.tp_q_head_num, layer.v_head_dim),
-            forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id),
-            forward_batch.token_to_kv_pool.get_value_buffer(layer.layer_id),
-            forward_batch.req_to_token_pool.req_to_token,
-            forward_batch.req_pool_indices,
-            forward_batch.seq_lens,
-            forward_batch.extend_seq_lens,
-            forward_batch.extend_start_loc,
-            max_extend_len,
-            layer.scaling,
-            layer.logit_cap,
-        )
-        return o
-
-    def forward_decode(
-        self,
-        q,
-        k,
-        v,
-        layer: RadixAttention,
-        forward_batch: ForwardBatch,
-        save_kv_cache=True,
-    ):
-        # During torch.compile, there is a bug in rotary_emb that causes the
-        # output value to have a 3D tensor shape. This reshapes the output correctly.
-        q = q.reshape(-1, layer.tp_q_head_num * layer.qk_head_dim)
-
-        # TODO: reuse the buffer across layers
-        if layer.qk_head_dim != layer.v_head_dim:
-            o = q.new_empty((q.shape[0], layer.tp_q_head_num * layer.v_head_dim))
-        else:
-            o = torch.empty_like(q)
-
-        # TODO: Add min seqlen
-        (
-            start_loc,
-            attn_logits,
-            max_seq_len,
-            min_seq_len,
-            max_extend_len,
-            ds_req_to_token,
-        ) = self.forward_metadata
-
-        k_label = torch.gather(
-            k,
-            2,
-            self.sorted_channels[layer.layer_id]
-            .unsqueeze(0)
-            .expand(k.shape[0], -1, -1),
-        )
-
-        if save_kv_cache:
-            forward_batch.token_to_kv_pool.set_kv_buffer(
-                layer, forward_batch.out_cache_loc, k, v, k_label
-            )
-
-        # NOTE(Andy) shouldn't be used when max_len_in_batch < heavy_token_num
-        #            and set a minimum value for sparse_decode
-        if (
-            min_seq_len < self.heavy_token_num
-            or max_seq_len < self.sparse_decode_threshold
-        ):
-            self.decode_attention_fwd(
-                q.view(-1, layer.tp_q_head_num, layer.qk_head_dim),
-                forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id),
-                forward_batch.token_to_kv_pool.get_value_buffer(layer.layer_id),
-                o.view(-1, layer.tp_q_head_num, layer.v_head_dim),
-                forward_batch.req_to_token_pool.req_to_token,
-                forward_batch.req_pool_indices,
-                start_loc,
-                forward_batch.seq_lens,
-                attn_logits,
-                max_seq_len,
-                layer.scaling,
-                layer.logit_cap,
-            )
-        else:
-            # TODO(Andy): indexing with torch.gather or torch.index_select or customized kernel
-            q_label = torch.gather(
-                q.view(-1, layer.tp_q_head_num, layer.qk_head_dim),
-                2,
-                self.sorted_channels[layer.layer_id]
-                .unsqueeze(0)
-                .expand(q.shape[0], -1, -1),
-            )
-            self.decode_sparse_attention_fwd(
-                q.view(-1, layer.tp_q_head_num, layer.qk_head_dim),
-                forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id),
-                forward_batch.token_to_kv_pool.get_value_buffer(layer.layer_id),
-                o.view(-1, layer.tp_q_head_num, layer.qk_head_dim),
-                q_label,
-                forward_batch.token_to_kv_pool.get_label_buffer(layer.layer_id),
-                ds_req_to_token,
-                forward_batch.seq_lens,
-                max_seq_len,
-                layer.scaling,
-                layer.logit_cap,
-                self.heavy_token_num,
-                self.att_out_approx,
-                self.mid_out,
-                self.mid_o_logexpsum,
-                self.BLOCK_SEQ,
-            )
-
-        return o

python/sglang/srt/mem_cache/memory_pool.py

@@ -1969,96 +1969,6 @@ class NSATokenToKVPool(MLATokenToKVPool):
         return kv_size_bytes
 
 
-class DoubleSparseTokenToKVPool(KVCache):
-    def __init__(
-        self,
-        size: int,
-        page_size: int,
-        dtype: torch.dtype,
-        head_num: int,
-        head_dim: int,
-        layer_num: int,
-        device: str,
-        heavy_channel_num: int,
-        enable_memory_saver: bool,
-        start_layer: Optional[int] = None,
-        end_layer: Optional[int] = None,
-    ):
-        super().__init__(
-            size,
-            page_size,
-            dtype,
-            layer_num,
-            device,
-            enable_memory_saver,
-            start_layer,
-            end_layer,
-        )
-
-        with self.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE):
-            with (
-                torch.cuda.use_mem_pool(self.custom_mem_pool)
-                if self.enable_custom_mem_pool
-                else nullcontext()
-            ):
-                # [size, head_num, head_dim] for each layer
-                self.k_buffer = [
-                    torch.zeros(
-                        (size + page_size, head_num, head_dim),
-                        dtype=dtype,
-                        device=device,
-                    )
-                    for _ in range(layer_num)
-                ]
-                self.v_buffer = [
-                    torch.zeros(
-                        (size + page_size, head_num, head_dim),
-                        dtype=dtype,
-                        device=device,
-                    )
-                    for _ in range(layer_num)
-                ]
-
-                # [size, head_num, heavy_channel_num] for each layer
-                self.label_buffer = [
-                    torch.zeros(
-                        (size + 1, head_num, heavy_channel_num),
-                        dtype=dtype,
-                        device=device,
-                    )
-                    for _ in range(layer_num)
-                ]
-
-    def get_key_buffer(self, layer_id: int):
-        return self.k_buffer[layer_id - self.start_layer]
-
-    def get_value_buffer(self, layer_id: int):
-        return self.v_buffer[layer_id - self.start_layer]
-
-    def get_label_buffer(self, layer_id: int):
-        return self.label_buffer[layer_id - self.start_layer]
-
-    def get_kv_buffer(self, layer_id: int):
-        return (
-            self.k_buffer[layer_id - self.start_layer],
-            self.v_buffer[layer_id - self.start_layer],
-        )
-
-    def set_kv_buffer(
-        self,
-        layer: RadixAttention,
-        loc: torch.Tensor,
-        cache_k: torch.Tensor,
-        cache_v: torch.Tensor,
-        cache_label: torch.Tensor,
-    ):
-        # NOTE(Andy): ignore the dtype check
-        layer_id = layer.layer_id
-        self.k_buffer[layer_id - self.start_layer][loc] = cache_k
-        self.v_buffer[layer_id - self.start_layer][loc] = cache_v
-        self.label_buffer[layer_id - self.start_layer][loc] = cache_label
-
-
 def move_kv_cache_native(
     k_buffer: List[torch.Tensor],
     v_buffer: List[torch.Tensor],

python/sglang/srt/model_executor/model_runner.py

@@ -18,7 +18,6 @@ from __future__ import annotations
 import datetime
 import gc
 import inspect
-import json
 import logging
 import os
 import socket
@@ -665,14 +664,6 @@ class ModelRunner(ModelRunnerKVCacheMixin):
                 # lora_manager.init_cuda_graph_batch_info().
                 self._init_lora_cuda_graph_moe_buffers()
 
-        # Init Double Sparsity
-        if server_args.enable_double_sparsity:
-            if server_args.ds_heavy_channel_type is None:
-                raise ValueError(
-                    "Please specify the heavy channel type for double sparsity optimization."
-                )
-            self.init_double_sparsity_channel_config(server_args.ds_heavy_channel_type)
-
         # Enable batch invariant mode
         if server_args.enable_deterministic_inference:
             from sglang.srt.batch_invariant_ops import enable_batch_invariant_mode
@@ -924,13 +915,6 @@ class ModelRunner(ModelRunnerKVCacheMixin):
     def model_specific_adjustment(self):
         server_args = self.server_args
 
-        if server_args.enable_double_sparsity:
-            logger.info(
-                "Double sparsity optimization is turned on. Use triton backend without CUDA graph."
-            )
-            server_args.attention_backend = "triton"
-            server_args.disable_cuda_graph = True
-
         if self.is_multimodal:
             if not self.is_multimodal_chunked_prefill_supported:
                 server_args.chunked_prefill_size = -1
@@ -2153,23 +2137,6 @@ class ModelRunner(ModelRunnerKVCacheMixin):
         full_attention_backend = ATTENTION_BACKENDS[backend_str](self)
         return attn_backend_wrapper(self, full_attention_backend)
 
-    def init_double_sparsity_channel_config(self, selected_channel):
-        selected_channel = "." + selected_channel + "_proj"
-        self.sorted_channels = []
-        # load channel config
-        with open(self.server_args.ds_channel_config_path, "r") as f:
-            channel_config = json.load(f)
-
-        for i in range(self.start_layer, self.end_layer):
-            key = "model.layers." + str(i) + ".self_attn" + selected_channel
-            self.sorted_channels.append(
-                torch.tensor(channel_config[key])[
-                    :, : self.server_args.ds_heavy_channel_num
-                ]
-                .contiguous()
-                .cuda()
-            )
-
     def kernel_warmup(self):
         """
         Warmup and tune kernels before cuda graph capture.

python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py

@@ -18,7 +18,6 @@ from sglang.srt.mem_cache.hisparse_memory_pool import (
     HiSparseTokenToKVPoolAllocator,
 )
 from sglang.srt.mem_cache.memory_pool import (
-    DoubleSparseTokenToKVPool,
     HybridLinearKVPool,
     HybridReqToTokenPool,
     MHATokenToKVPool,
@@ -409,20 +408,6 @@ class ModelRunnerKVCacheMixin:
                     start_layer=self.start_layer,
                     end_layer=self.end_layer,
                 )
-        elif self.server_args.enable_double_sparsity:
-            self.token_to_kv_pool = DoubleSparseTokenToKVPool(
-                self.max_total_num_tokens,
-                page_size=self.page_size,
-                dtype=self.kv_cache_dtype,
-                head_num=self.model_config.get_num_kv_heads(get_attention_tp_size()),
-                head_dim=self.model_config.head_dim,
-                layer_num=self.num_effective_layers,
-                device=self.device,
-                heavy_channel_num=self.server_args.ds_heavy_channel_num,
-                enable_memory_saver=self.server_args.enable_memory_saver,
-                start_layer=self.start_layer,
-                end_layer=self.end_layer,
-            )
         else:
             if self.is_hybrid_swa:
                 kwargs = {}

python/sglang/srt/server_args.py

@@ -591,14 +591,6 @@ class ServerArgs:
     dllm_algorithm: Optional[str] = None
     dllm_algorithm_config: Optional[str] = None
 
-    # Double Sparsity
-    enable_double_sparsity: bool = False
-    ds_channel_config_path: Optional[str] = None
-    ds_heavy_channel_num: int = 32
-    ds_heavy_token_num: int = 256
-    ds_heavy_channel_type: str = "qk"
-    ds_sparse_decode_threshold: int = 4096
-
     # Offloading
     cpu_offload_gb: int = 0
     offload_group_size: int = -1
@@ -5635,43 +5627,6 @@ class ServerArgs:
             help="The diffusion LLM algorithm configurations. Must be a YAML file.",
         )
 
-        # Double Sparsity
-        parser.add_argument(
-            "--enable-double-sparsity",
-            action="store_true",
-            help="Enable double sparsity attention",
-        )
-        parser.add_argument(
-            "--ds-channel-config-path",
-            type=str,
-            default=ServerArgs.ds_channel_config_path,
-            help="The path of the double sparsity channel config",
-        )
-        parser.add_argument(
-            "--ds-heavy-channel-num",
-            type=int,
-            default=ServerArgs.ds_heavy_channel_num,
-            help="The number of heavy channels in double sparsity attention",
-        )
-        parser.add_argument(
-            "--ds-heavy-token-num",
-            type=int,
-            default=ServerArgs.ds_heavy_token_num,
-            help="The number of heavy tokens in double sparsity attention",
-        )
-        parser.add_argument(
-            "--ds-heavy-channel-type",
-            type=str,
-            default=ServerArgs.ds_heavy_channel_type,
-            help="The type of heavy channels in double sparsity attention",
-        )
-        parser.add_argument(
-            "--ds-sparse-decode-threshold",
-            type=int,
-            default=ServerArgs.ds_sparse_decode_threshold,
-            help="The minimum decode sequence length required before the double-sparsity backend switches from the dense fallback to the sparse decode kernel.",
-        )
-
         # Offloading
         parser.add_argument(
             "--cpu-offload-gb",

test/manual/test_double_sparsity.py

@@ -1,65 +0,0 @@
-import os
-import unittest
-from types import SimpleNamespace
-
-from sglang.srt.utils import kill_process_tree
-from sglang.test.run_eval import run_eval
-from sglang.test.test_utils import (
-    DEFAULT_MODEL_NAME_FOR_TEST,
-    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
-    DEFAULT_URL_FOR_TEST,
-    CustomTestCase,
-    popen_launch_server,
-)
-
-
-class TestDoubleSparsity(CustomTestCase):
-    @classmethod
-    def setUpClass(cls):
-        cls.model = DEFAULT_MODEL_NAME_FOR_TEST
-        cls.base_url = DEFAULT_URL_FOR_TEST
-        dirpath = os.path.dirname(__file__)
-        config_file = os.path.join(
-            dirpath, "double-sparsity-config-Llama-3.1-8B-Instruct.json"
-        )
-        # NOTE: Generate the config file by running https://github.com/andy-yang-1/DoubleSparse/blob/main/evaluation/group_channel_config.py
-        cls.process = popen_launch_server(
-            cls.model,
-            cls.base_url,
-            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
-            other_args=[
-                "--enable-double-sparsity",
-                "--ds-channel-config-path",
-                config_file,
-                "--ds-heavy-channel-num",
-                "32",
-                "--ds-heavy-channel-type",
-                "k",
-                "--ds-heavy-token-num",
-                "512",
-                "--ds-sparse-decode-threshold",
-                "0",
-                "--max-total-tokens",
-                "200000",
-            ],
-        )
-
-    @classmethod
-    def tearDownClass(cls):
-        kill_process_tree(cls.process.pid)
-
-    def test_mmlu(self):
-        args = SimpleNamespace(
-            base_url=self.base_url,
-            model=self.model,
-            eval_name="mmlu",
-            num_examples=64,
-            num_threads=32,
-        )
-
-        metrics = run_eval(args)
-        self.assertGreaterEqual(metrics["score"], 0.65)
-
-
-if __name__ == "__main__":
-    unittest.main()