minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Remove deprecated double sparsity feature (#23009)

Browse Source
This commit is contained in:
Lianmin Zheng
2026-04-17 13:33:12 -07:00
committed by GitHub
parent 3df35ecc80
commit 44e67c6835
11 changed files with 2 additions and 1637 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/attention_registry.py
+2 -9
View File
@@ -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.double_sparsity_backend import (
DoubleSparseAttnBackend,
)
from sglang.srt.layers.attention.triton_backend import TritonAttnBackend
return DoubleSparseAttnBackend(runner)
else:
from sglang.srt.layers.attention.triton_backend import TritonAttnBackend
return TritonAttnBackend(runner)
return TritonAttnBackend(runner)
@register_attention_backend("torch_native")
python/sglang/srt/layers/attention/double_sparsity_backend.py
-257
View File
@@ -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/layers/attention/triton_ops/double_sparsity_attention.py
-1106
View File
File diff suppressed because it is too large Load Diff
python/sglang/srt/mem_cache/memory_pool.py
-90
View File
@@ -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
-33
View File
@@ -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
-15
View File
@@ -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
-45
View File
@@ -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",