From 3aa11ca722ce25d6d1f291636b3bcdd7b3f43dd9 Mon Sep 17 00:00:00 2001 From: Ke Bao Date: Tue, 6 Jan 2026 13:13:13 +0800 Subject: [PATCH] Remove hybrid_kvcache_ratio in server args (#16399) --- docs/basic_usage/llama4.md | 2 +- python/sglang/srt/configs/model_config.py | 43 +++++-------------- python/sglang/srt/managers/tp_worker.py | 2 +- .../model_runner_kv_cache_mixin.py | 28 +----------- python/sglang/srt/server_args.py | 12 +----- 5 files changed, 16 insertions(+), 71 deletions(-) diff --git a/docs/basic_usage/llama4.md b/docs/basic_usage/llama4.md index 1a2338a3f..05ffb2c60 100644 --- a/docs/basic_usage/llama4.md +++ b/docs/basic_usage/llama4.md @@ -32,7 +32,7 @@ python3 -m sglang.launch_server \ - **Chat Template**: Add `--chat-template llama-4` for chat completion tasks. - **Enable Multi-Modal**: Add `--enable-multimodal` for multi-modal capabilities. -- **Enable Hybrid-KVCache**: Add `--hybrid-kvcache-ratio` for hybrid kv cache. Details can be seen in [this PR](https://github.com/sgl-project/sglang/pull/6563) +- **Enable Hybrid-KVCache**: Set `--swa-full-tokens-ratio` to adjust the ratio of SWA layer (for Llama4, it's local attention layer) KV tokens / full layer KV tokens. (default: 0.8, range: 0-1) ### EAGLE Speculative Decoding diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index 78aaf5c65..aa10cb08d 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -94,9 +94,6 @@ class ModelConfig: quantization: Optional[str] = None, override_config_file: Optional[str] = None, is_draft_model: bool = False, - hybrid_kvcache_ratio: Optional[ - float - ] = None, # TODO: remove this, it is not a model config model_impl: Union[str, ModelImpl] = ModelImpl.AUTO, sampling_defaults: str = "openai", quantize_and_serve: bool = False, @@ -199,7 +196,7 @@ class ModelConfig: self._derive_model_shapes() # Update hybrid model - self._derive_hybrid_model(hybrid_kvcache_ratio) + self._derive_hybrid_model() # Verify quantization self._verify_quantization() @@ -251,7 +248,6 @@ class ModelConfig: enable_multimodal=server_args.enable_multimodal, dtype=server_args.dtype, quantization=quantization, - hybrid_kvcache_ratio=server_args.hybrid_kvcache_ratio, model_impl=server_args.model_impl, sampling_defaults=server_args.sampling_defaults, quantize_and_serve=server_args.quantize_and_serve, @@ -304,15 +300,11 @@ class ModelConfig: self.hf_config.architectures[0] = "Qwen3NextForCausalLMMTP" self.hf_config.num_nextn_predict_layers = 1 - def _derive_hybrid_model(self, hybrid_kvcache_ratio: Optional[float] = None): + def _derive_hybrid_model(self): # Use self.context_len after it has been initialized to prevent using context_len which may be None. - self.is_hybrid_swa = is_hybrid_model( - self.hf_config.architectures, - hybrid_kvcache_ratio=hybrid_kvcache_ratio, - context_length=self.context_len, - attention_chunk_size=self.attention_chunk_size, - ) - if self.is_hybrid_swa is not None: + self.is_hybrid_swa = is_hybrid_swa_model(self.hf_config.architectures) + + if self.is_hybrid_swa: self.swa_attention_layer_ids, self.full_attention_layer_ids = ( get_hybrid_layer_ids( self.hf_config.architectures, @@ -1151,27 +1143,14 @@ def yarn_get_mscale(scale: float = 1, mscale: float = 1) -> float: return 0.1 * mscale * math.log(scale) + 1.0 -def is_hybrid_model( - model_architectures: List[str], - hybrid_kvcache_ratio: Optional[float], - context_length: Optional[int], - attention_chunk_size: Optional[int], -): - if model_architectures[0] in [ +def is_hybrid_swa_model(model_architectures: List[str]): + + hybrid_swa_archs = { + "Llama4ForConditionalGeneration", "MiMoV2FlashForCausalLM", "MiMoV2MTP", - ]: - return 1 - if hybrid_kvcache_ratio is None: - return None - elif ( - hybrid_kvcache_ratio > 0 - and model_architectures[0] == "Llama4ForConditionalGeneration" - and context_length > attention_chunk_size - ): - return hybrid_kvcache_ratio - else: - return None + } + return any(arch in hybrid_swa_archs for arch in model_architectures) def get_hybrid_layer_ids( diff --git a/python/sglang/srt/managers/tp_worker.py b/python/sglang/srt/managers/tp_worker.py index 1f1875254..49f63a198 100644 --- a/python/sglang/srt/managers/tp_worker.py +++ b/python/sglang/srt/managers/tp_worker.py @@ -71,7 +71,7 @@ class BaseTpWorker(ABC): @property def is_hybrid_swa(self) -> bool: - return self.model_runner.is_hybrid_swa is not None + return self.model_runner.is_hybrid_swa def get_tokens_per_layer_info(self): return ( 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 98559e226..5cd65ab16 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 @@ -125,8 +125,6 @@ class ModelRunnerKVCacheMixin: ) elif mambaish := self.mambaish_config: num_layers = len(mambaish.full_attention_layer_ids) - elif self.model_config.full_attention_layer_ids: - num_layers = len(self.model_config.full_attention_layer_ids) else: num_layers = self.num_effective_layers @@ -202,31 +200,7 @@ class ModelRunnerKVCacheMixin: def set_num_tokens_hybrid_swa(self: ModelRunner): page_size = self.server_args.page_size - if ( - "Llama4ForConditionalGeneration" - in self.model_config.hf_config.architectures - ): - temp_ratio = ( - (1 - self.is_hybrid_swa) - + self.is_hybrid_swa - * self.attention_chunk_size - / self.model_config.context_len - ) - self.swa_max_total_num_tokens = ( - 4 * self.max_total_num_tokens * temp_ratio // (3 * temp_ratio + 1) - ) - self.full_max_total_num_tokens = ( - 4 * self.max_total_num_tokens - - 12 * self.max_total_num_tokens * temp_ratio // (3 * temp_ratio + 1) - ) - self.swa_max_total_num_tokens = ( - self.swa_max_total_num_tokens // page_size * page_size - ) - self.full_max_total_num_tokens = ( - self.full_max_total_num_tokens // page_size * page_size - ) - self.max_total_num_tokens = self.full_max_total_num_tokens - elif "MiMoV2MTP" in self.model_config.hf_config.architectures: + if "MiMoV2MTP" in self.model_config.hf_config.architectures: assert self.is_draft_worker # MiMoV2MTP uses SWA, so set full KV cache to 0 self.full_max_total_num_tokens = 0 diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index b92ec00d3..a2b26e0e0 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -308,7 +308,6 @@ class ServerArgs: priority_scheduling_preemption_threshold: int = 10 schedule_conservativeness: float = 1.0 page_size: Optional[int] = None - hybrid_kvcache_ratio: Optional[float] = None swa_full_tokens_ratio: float = 0.8 disable_hybrid_swa_memory: bool = False radix_eviction_policy: str = "lru" @@ -2849,15 +2848,8 @@ class ServerArgs: ) parser.add_argument( "--hybrid-kvcache-ratio", - nargs="?", - const=0.5, - type=float, - default=ServerArgs.hybrid_kvcache_ratio, - help=( - "Mix ratio in [0,1] between uniform and hybrid kv buffers " - "(0.0 = pure uniform: swa_size / full_size = 1)" - "(1.0 = pure hybrid: swa_size / full_size = local_attention_size / context_length)" - ), + action=DeprecatedAction, + help="Note: --hybrid-kvcache-ratio is deprecated now. Please use --swa-full-tokens-ratio instead.", ) parser.add_argument( "--swa-full-tokens-ratio",