fix(gemma4): only apply swa_full_tokens_ratio=0.15 to MoE variants

python/sglang/srt/server_args.py

@@ -2240,40 +2240,62 @@ def _handle_model_specific_adjustments(self):
                         "SM100 for Gemma-4 (modelopt_fp4)"
                     )
 
-            # Gemma-4 uses a 25:5 sliding-window : full-attention layer ratio
-            # (see ``Gemma4TextConfig.layer_types``).  The shipped default
+            # Gemma-4 uses a 5:1 SWA:full-attention layer ratio (see
+            # ``Gemma4TextConfig.layer_types``).  The shipped default
             # ``swa_full_tokens_ratio = 0.8`` is tuned for models where the
-            # sliding-window pool is the binding constraint, but for Gemma-4
-            # the full-attention pool is binding under concurrent long-context
-            # workloads: with the default ratio the full pool only fits ~65
-            # 9k-token requests on a 180 GB B200, forcing partial KV eviction
-            # and re-prefill (visible as ``#cached-token: 1003 #new-token:
-            # 7010`` lines in the serving log) under typical 80-request
-            # summarization loads.
+            # sliding-window pool is the binding constraint, but for the
+            # **MoE** Gemma-4 (``26B-A4B-IT``: 30 layers = 25 SWA + 5 full,
+            # 128 experts top-k 8) the full-attention pool is binding under
+            # concurrent long-context workloads.  Lowering the ratio to
+            # ``0.15`` shifts memory from the over-provisioned SWA pool to
+            # the under-provisioned full pool; median summarization TTFT
+            # drops 16% (10.5 s -> 8.7 s) on B200 with no MMLU regression.
             #
-            # Lowering the ratio to ~0.15 shifts memory from the over-
-            # provisioned SWA pool (25 layers × 1024-token window) to the
-            # under-provisioned full pool (5 layers × full context length).
-            # On the same 180 GB B200, the full pool grows from ~594 k tokens
-            # to ~2.14 M tokens (3.6× larger; enough for ~237 concurrent
-            # 9k-token requests), while the SWA pool shrinks from ~475 k to
-            # ~321 k tokens (still ~313 concurrent 1024-token windows,
-            # far above any realistic request count).  Median TTFT on a
-            # summarization workload of 80 × 8k-input / 1k-output prompts
-            # drops 16.5 % (10.5 s -> 8.7 s) on a B200 with TP=1, MTP, and
-            # the triton attention backend, with no MMLU regression.
+            # **Do not apply** this override to dense Gemma-4 variants
+            # (``31B-it``, ``E4B-IT``) — they have less GPU memory free
+            # after model load (dense weights take more RAM than MoE
+            # sparse weights), so the SWA pool becomes critically small
+            # at this ratio and chokes admission under high concurrency.
+            # Empirically: applying ``0.15`` to 31B on B200 with 80
+            # concurrent 1k/1k chat requests caused SWA usage to hit
+            # 100% saturation and dropped output throughput by ~3x.
             #
-            # Only apply when the user did not explicitly set the ratio,
-            # mirroring the pattern in ``apply_deepseek_v4_defaults``.
-            if self.swa_full_tokens_ratio == ServerArgs.swa_full_tokens_ratio:
+            # MoE detection via ``num_experts`` on the text config — same
+            # pattern used in ``gemma4_causal.py:1166``.  Also keep the
+            # ``apply_deepseek_v4_defaults``-style "respect user override"
+            # predicate (note: the predicate currently can't distinguish
+            # user-passed ``0.8`` from the dataclass default; same caveat
+            # as the upstream DSV4 override).
+            try:
+                _hf_text_config = self.get_model_config().hf_text_config
+            except Exception:
+                _hf_text_config = None
+            _gemma4_num_experts = (
+                int(getattr(_hf_text_config, "num_experts", 0) or 0)
+                if _hf_text_config is not None
+                else 0
+            )
+            _is_gemma4_moe = _gemma4_num_experts > 0
+            if (
+                _is_gemma4_moe
+                and self.swa_full_tokens_ratio == ServerArgs.swa_full_tokens_ratio
+            ):
                 self.swa_full_tokens_ratio = 0.15
                 logger.info(
                     "Setting swa_full_tokens_ratio to "
                     f"{self.swa_full_tokens_ratio} for {model_arch} "
-                    "(Gemma-4 has a 25:5 SWA:full layer split; the default "
-                    "ratio over-provisions the SWA pool and under-provisions "
-                    "the full-attention pool, causing partial KV eviction "
-                    "and re-prefill under concurrent long-context loads)."
+                    f"(MoE Gemma-4 with num_experts={_gemma4_num_experts}; "
+                    "the default ratio over-provisions the SWA pool and "
+                    "under-provisions the full-attention pool, causing "
+                    "partial KV eviction and re-prefill under concurrent "
+                    "long-context loads)."
+                )
+            elif not _is_gemma4_moe:
+                logger.info(
+                    f"Keeping default swa_full_tokens_ratio="
+                    f"{self.swa_full_tokens_ratio} for {model_arch} "
+                    "(dense Gemma-4; MoE-specific 0.15 override skipped "
+                    "to avoid SWA pool starvation)."
                 )
         elif model_arch == "MossVLForConditionalGeneration":
             if self.is_attention_backend_not_set():

test/srt/test_gemma4_swa_full_tokens_ratio.py

@@ -46,13 +46,19 @@ def _stub_sm100(monkeypatch):
     monkeypatch.setattr(srv_args, "is_sm100_supported", lambda: True, raising=False)
 
 
-def _invoke_gemma4_adjustment(args, model_arch="Gemma4ForCausalLM"):
+def _invoke_gemma4_adjustment(
+    args, model_arch="Gemma4ForCausalLM", num_experts=0
+):
     """Run only the small Gemma-4 branch of ``_handle_model_specific_adjustments``.
 
     The full method walks every supported model family and pulls in lots of
     HF-config-touching helpers; we copy just the Gemma-4 logic that exercises
     the SWA override under test.  Keeping the test scope tight avoids
     coupling it to unrelated branches.
+
+    ``num_experts`` simulates ``hf_text_config.num_experts`` so we can
+    cover both MoE Gemma-4 (26B-A4B-IT, ``num_experts=128``) and dense
+    Gemma-4 (31B-it / E4B-IT, ``num_experts=0``).
     """
     from sglang.srt.server_args import ServerArgs
 
@@ -62,53 +68,74 @@ def _invoke_gemma4_adjustment(args, model_arch="Gemma4ForCausalLM"):
         "Gemma4ForConditionalGeneration",
         "Gemma4ForCausalLM",
     )
-    if args.swa_full_tokens_ratio == ServerArgs.swa_full_tokens_ratio:
+    # Mirror the MoE-only gating logic from server_args.py.
+    _is_gemma4_moe = num_experts > 0
+    if (
+        _is_gemma4_moe
+        and args.swa_full_tokens_ratio == ServerArgs.swa_full_tokens_ratio
+    ):
         args.swa_full_tokens_ratio = 0.15
 
 
-def test_default_overridden_for_gemma4():
-    """Unset ratio should be overridden to 0.15 for Gemma-4."""
+def test_moe_gemma4_default_overridden():
+    """MoE Gemma-4 (e.g. 26B-A4B-IT) should get the 0.15 override when ratio is unset."""
     args = _make_args()
     assert args.swa_full_tokens_ratio == ServerArgs.swa_full_tokens_ratio  # default 0.8
-    _invoke_gemma4_adjustment(args)
+    _invoke_gemma4_adjustment(args, num_experts=128)  # 26B-A4B-IT has 128 experts
     assert args.swa_full_tokens_ratio == 0.15
 
 
+def test_dense_gemma4_default_preserved():
+    """Dense Gemma-4 (e.g. 31B-it, E4B-IT) should KEEP the upstream default 0.8.
+
+    Applying 0.15 to dense variants causes SWA pool starvation under high
+    concurrency (verified on 31B + B200: SWA hits 100% saturation,
+    output throughput collapses by ~3x).  See
+    ``agent-pad/runs/.../benchmark_final/FINAL_COMPARISON.md``.
+    """
+    args = _make_args()
+    expected = ServerArgs.swa_full_tokens_ratio  # 0.8
+    _invoke_gemma4_adjustment(args, num_experts=0)  # dense
+    assert args.swa_full_tokens_ratio == expected
+
+
 @pytest.mark.parametrize(
     "model_arch", ["Gemma4ForCausalLM", "Gemma4ForConditionalGeneration"]
 )
 def test_user_override_preserved(model_arch):
-    """If user passes --swa-full-tokens-ratio, it must be respected."""
+    """If user passes --swa-full-tokens-ratio, it must be respected (MoE case)."""
     args = _make_args(swa_full_tokens_ratio=0.5)
-    _invoke_gemma4_adjustment(args, model_arch)
+    _invoke_gemma4_adjustment(args, model_arch, num_experts=128)
     assert args.swa_full_tokens_ratio == 0.5
 
     args = _make_args(swa_full_tokens_ratio=1.0)
-    _invoke_gemma4_adjustment(args, model_arch)
+    _invoke_gemma4_adjustment(args, model_arch, num_experts=128)
     assert args.swa_full_tokens_ratio == 1.0
 
 
-def test_full_method_runs_for_gemma4_for_causal_lm(monkeypatch):
-    """Smoke test: invoke the real ``_handle_model_specific_adjustments`` and
-    assert the SWA ratio path fires alongside the attention-backend setup.
+def test_full_method_runs_for_moe_gemma4(monkeypatch):
+    """Smoke test for MoE Gemma-4: invoke the real
+    ``_handle_model_specific_adjustments`` and assert the SWA ratio path
+    fires alongside the attention-backend setup.
 
     We stub the model-config loader so we don't need real Gemma-4 weights.
     """
     from sglang.srt.server_args import ServerArgs
 
     args = _make_args(
-        model_path="fake-gemma4",
+        model_path="fake-gemma4-moe",
         attention_backend=None,
         prefill_attention_backend=None,
         decode_attention_backend=None,
         moe_runner_backend="auto",
     )
 
-    # ``_handle_model_specific_adjustments`` resolves ``model_arch`` from
-    # ``self.get_model_config()``; stub that to return our synthetic Gemma-4.
+    class _FakeTextConfig:
+        num_experts = 128
+
     class _FakeModelConfig:
         quantization = None
-        hf_text_config = None
+        hf_text_config = _FakeTextConfig()
 
     class _FakeModelArchConfig:
         def __init__(self):
@@ -138,5 +165,54 @@ def _fake_get_model_config(self):
     assert args.attention_backend in ("triton", "trtllm_mha")
 
 
+def test_full_method_runs_for_dense_gemma4(monkeypatch):
+    """Smoke test for dense Gemma-4: invoke the real method and assert
+    the override is SKIPPED (default 0.8 preserved)."""
+    from sglang.srt.server_args import ServerArgs
+
+    args = _make_args(
+        model_path="fake-gemma4-dense",
+        attention_backend=None,
+        prefill_attention_backend=None,
+        decode_attention_backend=None,
+        moe_runner_backend="auto",
+    )
+
+    class _FakeTextConfig:
+        num_experts = 0  # dense (or attribute missing → also evaluates to 0)
+
+    class _FakeModelConfig:
+        quantization = None
+        hf_text_config = _FakeTextConfig()
+
+    class _FakeModelArchConfig:
+        def __init__(self):
+            self.architectures = ["Gemma4ForCausalLM"]
+
+    def _fake_get_model_arch_config(self):
+        return _FakeModelArchConfig()
+
+    def _fake_get_model_config(self):
+        return _FakeModelConfig()
+
+    monkeypatch.setattr(
+        ServerArgs, "get_model_arch_config", _fake_get_model_arch_config, raising=False
+    )
+    monkeypatch.setattr(
+        ServerArgs, "get_model_config", _fake_get_model_config, raising=False
+    )
+
+    try:
+        args._handle_model_specific_adjustments()
+    except Exception as exc:
+        pytest.skip(
+            f"_handle_model_specific_adjustments needs more stubs in this env: {exc}"
+        )
+
+    # Dense Gemma-4: override should NOT fire, ratio stays at upstream default 0.8.
+    assert args.swa_full_tokens_ratio == ServerArgs.swa_full_tokens_ratio
+    assert args.attention_backend in ("triton", "trtllm_mha")
+
+
 if __name__ == "__main__":
     raise SystemExit(pytest.main([__file__, "-v"]))