Disable swa memory for trtllm-mha backend (#17429)

python/sglang/srt/configs/model_config.py

@@ -101,6 +101,7 @@ class ModelConfig:
         is_multi_layer_eagle: bool = False,
         encoder_only: bool = False,
         language_only: bool = False,
+        disable_hybrid_swa_memory: bool = False,
     ) -> None:
         # Parse args
         self.model_path = model_path
@@ -111,6 +112,7 @@ class ModelConfig:
         self.sampling_defaults = sampling_defaults
         self.quantize_and_serve = quantize_and_serve
         self.is_multi_layer_eagle = is_multi_layer_eagle
+        self.disable_hybrid_swa_memory = disable_hybrid_swa_memory
 
         # Validate quantize_and_serve configuration
         self._validate_quantize_and_serve_config()
@@ -257,6 +259,7 @@ class ModelConfig:
             language_only=server_args.language_only,
             encoder_only=server_args.encoder_only,
             is_draft_model=is_draft_model,
+            disable_hybrid_swa_memory=server_args.disable_hybrid_swa_memory,
             **kwargs,
         )
 
@@ -310,7 +313,10 @@ 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)
+        self.is_hybrid_swa = (
+            is_hybrid_swa_model(self.hf_config.architectures)
+            and not self.disable_hybrid_swa_memory
+        )
 
         if self.is_hybrid_swa:
             self.swa_attention_layer_ids, self.full_attention_layer_ids = (

python/sglang/srt/server_args.py

@@ -1280,9 +1280,9 @@ class ServerArgs:
                 or decode_attn_backend == "trtllm_mha"
             ):
                 # TODO: support swa kv indices translation for trtllm_mha attention backend
-                self.swa_full_tokens_ratio = 1.0
+                self.disable_hybrid_swa_memory = True
                 logger.warning(
-                    "Set swa_full_tokens_ratio to 1.0 for GPT-OSS model with trtllm_mha attention backend."
+                    "Disable hybrid SWA memory for GPT-OSS model with trtllm_mha attention backend."
                 )
 
             quant_method = get_quantization_config(hf_config)