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

SWA Prefix Cache (#7367)

Browse Source 
Co-authored-by: Ying Sheng <sqy1415@gmail.com>
This commit is contained in:
Hanming Lu
2025-07-13 12:31:07 -07:00
committed by GitHub
parent 0c55cbcfc5
commit 9379da77de
16 changed files with 1742 additions and 158 deletions
Download Patch File Download Diff File Expand all files Collapse all files

77
python/sglang/srt/model_executor/model_runner.py
View File

@@ -275,6 +275,15 @@ class ModelRunner:
self.sampler = Sampler()
self.load_model()
if (
not self.server_args.disable_hybrid_swa_memory
and self.sliding_window_size is not None
and self.sliding_window_size > 0
):
architectures = self.model_config.hf_config.architectures
if architectures and not any("Llama4" in arch for arch in architectures):
self.is_hybrid = self.model_config.is_hybrid = True
self.start_layer = getattr(self.model, "start_layer", 0)
self.end_layer = getattr(
self.model, "end_layer", self.model_config.num_hidden_layers
@@ -471,10 +480,6 @@ class ModelRunner:
if self.model_config.context_len > 8192:
self.mem_fraction_static *= 0.85
if self.is_hybrid and not server_args.disable_radix_cache:
logger.info("Automatically disable radix cache for hybrid cache.")
server_args.disable_radix_cache = True
def init_torch_distributed(self):
logger.info("Init torch distributed begin.")
@@ -645,11 +650,15 @@ class ModelRunner:
)
# Parse other args
self.sliding_window_size = (
self.model.get_attention_sliding_window_size()
if hasattr(self.model, "get_attention_sliding_window_size")
else None
)
self.sliding_window_size = None
if hasattr(self.model, "get_attention_sliding_window_size"):
self.sliding_window_size = self.model.get_attention_sliding_window_size()
elif self.model_config.attention_chunk_size is not None:
self.sliding_window_size = self.model_config.attention_chunk_size
print(
f"Setting sliding_window_size to be attention_chunk_size: {self.sliding_window_size}"
)
self.dtype = self.model_config.dtype
after_avail_memory = get_available_gpu_memory(self.device, self.gpu_id)
@@ -992,8 +1001,53 @@ class ModelRunner:
)
self.max_total_num_tokens = self.full_max_total_num_tokens
else:
raise ValueError(
f"Unsupported model for hybrid cache: {self.model_config.hf_config.architectures}."
assert self.sliding_window_size is not None and self.sliding_window_size > 0
full_attention_layer_ids = []
swa_attention_layer_ids = []
try:
layers = self.model.model.layers
except:
try:
layers = self.model.language_model.model.layers
except:
self.is_hybrid = False
return
for layer in layers:
if (
layer.self_attn.attn.sliding_window_size is None
or layer.self_attn.attn.sliding_window_size == -1
):
full_attention_layer_ids.append(layer.layer_id)
else:
swa_attention_layer_ids.append(layer.layer_id)
self.model_config.swa_attention_layer_ids = swa_attention_layer_ids
self.model_config.full_attention_layer_ids = full_attention_layer_ids
# Algorithm:
# Existing max_total_num_tokens is per layer and assume all layers have the same number of tokens.
# - Find total # of tokens available across layers.
# - Calculate full_max_total_num_tokens and swa_max_total_num_tokens based on the given swa_full_tokens_ratio.
total_tokens = (
self.max_total_num_tokens * self.model_config.num_hidden_layers
)
full_layers_num = len(full_attention_layer_ids)
swa_layers_num = len(swa_attention_layer_ids)
swa_full_tokens_ratio = self.server_args.swa_full_tokens_ratio
# Solve the equations:
# 1. swa_max_total_num_tokens * swa_layers_num + full_max_total_num_tokens * full_layers_num == total_tokens
# 2. full_max_total_num_tokens * swa_full_tokens_ratio == swa_max_total_num_tokens
denominator = swa_full_tokens_ratio * swa_layers_num + full_layers_num
self.full_max_total_num_tokens = int(total_tokens / denominator)
self.swa_max_total_num_tokens = int(
self.full_max_total_num_tokens * swa_full_tokens_ratio
)
self.max_total_num_tokens = self.full_max_total_num_tokens
logger.info(
f"Use Sliding window memory pool. full_layer_tokens={self.full_max_total_num_tokens}, swa_layer_tokens={self.swa_max_total_num_tokens}"
)
def init_memory_pool(
@@ -1072,7 +1126,6 @@ class ModelRunner:
// self.server_args.page_size
* self.server_args.page_size
)
# create token size for hybrid cache
if self.is_hybrid:
self.set_num_token_hybrid()