diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index e3a6035b0..a3153bac1 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -504,7 +504,7 @@ class ModelConfig: 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_swa_model(self.hf_config.architectures) + is_hybrid_swa_model(self.hf_config.architectures, self.hf_text_config) and not self.disable_hybrid_swa_memory ) @@ -1671,7 +1671,10 @@ def compute_mla_mscale_scaling(rope_scaling: dict, base_scaling: float) -> float return base_scaling * mscale * mscale -def is_hybrid_swa_model(model_architectures: List[str]): +def is_hybrid_swa_model( + model_architectures: List[str], + hf_text_config: Optional[PretrainedConfig] = None, +): hybrid_swa_archs = { "Llama4ForConditionalGeneration", @@ -1687,7 +1690,13 @@ def is_hybrid_swa_model(model_architectures: List[str]): "Gemma4ForConditionalGeneration", "LagunaForCausalLM", } - return any(arch in hybrid_swa_archs for arch in model_architectures) + if any(arch in hybrid_swa_archs for arch in model_architectures): + return True + # Also recognize models that explicitly opt-in via their HF text config, + # so custom hybrid-SWA architectures don't need to be added to the allowlist. + if hf_text_config is not None and getattr(hf_text_config, "is_hybrid_swa", False): + return True + return False def get_hybrid_layer_ids( @@ -1758,6 +1767,17 @@ def get_hybrid_layer_ids( full_attention_layer_ids = [ i for i, x in enumerate(layer_types) if x == "full_attention" ] + elif getattr(hf_text_config, "hybrid_layer_pattern", None) is not None: + # Generic fallback for custom hybrid SWA models that opt in via + # hf_text_config.is_hybrid_swa and expose a hybrid_layer_pattern + # (1 = SWA, 0 = full) without needing to be added to the allowlist. + hybrid_layer_pattern = hf_text_config.hybrid_layer_pattern + swa_attention_layer_ids = [ + i for i in range(num_hidden_layers) if hybrid_layer_pattern[i] == 1 + ] + full_attention_layer_ids = [ + i for i in range(num_hidden_layers) if hybrid_layer_pattern[i] == 0 + ] else: swa_attention_layer_ids = None full_attention_layer_ids = None