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

Fix swa kv cache memory allocation (#18039)

Browse Source
This commit is contained in:
Ke Bao
2026-02-01 14:26:51 +08:00
committed by GitHub
parent 38d275a9fd
commit d396650bd2
2 changed files with 87 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
python/sglang/srt/configs/model_config.py
View File

@@ -392,6 +392,17 @@ class ModelConfig:
self.head_dim,
)
self.swa_head_dim = getattr(
self.hf_text_config,
"swa_head_dim",
self.head_dim,
)
self.swa_v_head_dim = getattr(
self.hf_text_config,
"swa_v_head_dim",
self.v_head_dim,
)
# FIXME: temporary special judge for MLA architecture
if (
"DeepseekV2ForCausalLM" in self.hf_config.architectures
@@ -592,12 +603,11 @@ class ModelConfig:
def get_swa_num_kv_heads(self, tensor_parallel_size) -> int:
"""Similar to get_num_kv_heads(), but for SWA."""
if not self.is_hybrid_swa_compress:
return 0
# For MiMoV2FlashForCausalLM models
total_num_kv_heads = self.hf_text_config.swa_num_key_value_heads
return max(1, total_num_kv_heads // tensor_parallel_size)
if hasattr(self.hf_text_config, "swa_num_key_value_heads"):
total_num_kv_heads = self.hf_text_config.swa_num_key_value_heads
return max(1, total_num_kv_heads // tensor_parallel_size)
else:
return self.get_num_kv_heads(tensor_parallel_size)
# adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/config.py
def _parse_quant_hf_config(self):

99
python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py
View File

@@ -79,12 +79,28 @@ class ModelRunnerKVCacheMixin:
)
cell_size += indexer_size_per_token * num_layers * element_size
else:
cell_size = (
self.model_config.get_num_kv_heads(get_attention_tp_size())
* (self.model_config.head_dim + self.model_config.v_head_dim)
* num_layers
* kv_size
)
if self.model_config.is_hybrid_swa:
full_layers_num = len(self.model_config.full_attention_layer_ids)
swa_layers_num = len(self.model_config.swa_attention_layer_ids)
full_per_token = self.model_config.get_num_kv_heads(
get_attention_tp_size()
) * (self.model_config.head_dim + self.model_config.v_head_dim)
swa_per_token = self.model_config.get_swa_num_kv_heads(
get_attention_tp_size()
) * (self.model_config.swa_head_dim + self.model_config.swa_v_head_dim)
cell_size = (
full_per_token * full_layers_num + swa_per_token * swa_layers_num
) * kv_size
else:
cell_size = (
self.model_config.get_num_kv_heads(get_attention_tp_size())
* (self.model_config.head_dim + self.model_config.v_head_dim)
* num_layers
* kv_size
)
if is_float4_e2m1fn_x2(self.kv_cache_dtype):
# kv_scale_buffer
@@ -95,17 +111,6 @@ class ModelRunnerKVCacheMixin:
cell_size = (cell_size // 2) + (
(n * k * num_layers * 2 * kv_size) // scale_block_size
)
if "MiMoV2FlashForCausalLM" in self.model_config.hf_config.architectures:
cell_size += (
self.model_config.get_swa_num_kv_heads(get_attention_tp_size())
* (
self.model_config.hf_text_config.swa_head_dim
+ self.model_config.hf_text_config.swa_v_head_dim
)
* len(self.model_config.swa_attention_layer_ids)
* kv_size
)
return cell_size
def profile_max_num_token(self: ModelRunner, total_gpu_memory: int):
@@ -225,21 +230,59 @@ class ModelRunnerKVCacheMixin:
)
return
# 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
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
# Use unified memory-based allocation for all hybrid SWA models.
#
# Let:
# F = Full layer per-token memory
# S = SWA layer per-token memory (may differ from F)
# r = swa_full_tokens_ratio = swa_tokens / full_tokens
#
# The profile phase computed:
# cell_size = F * n_full + S * n_swa
# max_total_num_tokens = rest_memory / cell_size
# => total_memory = max_total_num_tokens * (F * n_full + S * n_swa)
#
# We need to solve:
# full_tokens * F * n_full + swa_tokens * S * n_swa = total_memory
# swa_tokens = full_tokens * r
#
# Solution:
# full_tokens = total_memory / (F * n_full + r * S * n_swa)
# = max_total_num_tokens * (F * n_full + S * n_swa) / (F * n_full + r * S * n_swa)
kv_size = torch._utils._element_size(self.kv_cache_dtype)
# Full layer per-token memory
full_per_token = (
self.model_config.get_num_kv_heads(get_attention_tp_size())
* (self.model_config.head_dim + self.model_config.v_head_dim)
* kv_size
)
# SWA layer per-token memory
swa_per_token = (
self.model_config.get_swa_num_kv_heads(get_attention_tp_size())
* (self.model_config.swa_head_dim + self.model_config.swa_v_head_dim)
* kv_size
)
# Total memory available from profile
total_memory = self.max_total_num_tokens * (
full_per_token * full_layers_num + swa_per_token * swa_layers_num
)
# Solve the equations
denominator = (
full_per_token * full_layers_num
+ swa_full_tokens_ratio * swa_per_token * swa_layers_num
)
assert (
denominator > 0
), f"Invalid denominator={denominator} for swa_full_tokens_ratio={swa_full_tokens_ratio} and swa_layers_num={swa_layers_num} and full_layers_num={full_layers_num}"
self.full_max_total_num_tokens = int(total_tokens / denominator)
), f"Invalid denominator={denominator} for memory-based allocation. full_per_token={full_per_token}, full_layers_num={full_layers_num}, swa_per_token={swa_per_token}, swa_layers_num={swa_layers_num}, swa_full_tokens_ratio={swa_full_tokens_ratio}"
self.full_max_total_num_tokens = int(total_memory / denominator)
self.swa_max_total_num_tokens = int(
self.full_max_total_num_tokens * swa_full_tokens_ratio
)