diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index b0fbc87dc..43543b655 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -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): diff --git a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py index 6cfa91e87..c7dcd5766 100644 --- a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py +++ b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py @@ -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 )