From b0e87f3a5cca986617be46a1dac00f46d8251c2b Mon Sep 17 00:00:00 2001 From: Pengyu Chen Date: Sat, 23 May 2026 08:57:03 +0000 Subject: [PATCH] fix(gemma4): only apply swa_full_tokens_ratio=0.15 to MoE variants Patch 2 (PR pyc96/sglang#2) set swa_full_tokens_ratio=0.15 for every Gemma-4 model. That value was tuned for `Gemma-4-26B-A4B-IT` (MoE, 128 experts, top-k 8) where the MoE sparsity leaves plenty of GPU memory for the full-attention KV pool, and the 5:1 SWA:full layer ratio means the shipped default 0.8 over-provisions the SWA pool. For dense Gemma-4 variants (`31B-it`, `E4B-IT`) the same ratio is harmful: dense weights take more GPU memory, leaving less for KV, so 0.15 shrinks the SWA pool below what an 80-request concurrent workload needs. Empirically (on `gemma-4-31B-it` + trtllm_mha + MTP + 1x B200 with 80 concurrent 1k/1k chat requests): ratio=0.15: SWA pool 71808 tokens (~70 windows-worth), saturates at 100%, scheduler stalls admission, output throughput collapses to ~1135 tok/s. ratio=0.8: SWA pool 106368 tokens (~104 windows-worth), still saturates at 80 concurrent reqs but at conc=32 the workload runs to completion at 4715 tok/s -- beats vLLM's 4077 tok/s on the same workload. This commit gates the 0.15 override on `num_experts > 0`, read from the model's `hf_text_config`. Mirrors the MoE-detection pattern in `gemma4_causal.py:1166`. Per-model verification on 1x B200: 26B-A4B-IT (MoE, num_experts=128): log: 'Setting swa_full_tokens_ratio to 0.15 for ... ' pool: full_layer_tokens=2138240 swa_layer_tokens=320704 (unchanged from Patch 2 -- regression-safe) 31B-it (dense, num_experts=0): log: 'Keeping default swa_full_tokens_ratio=0.8 ... ' pool: full_layer_tokens=132992 swa_layer_tokens=106368 (instead of the broken 478720 / 71808 layout from Patch 2) E4B-IT (dense, num_experts=0): same MoE-only-skipped path as 31B. Benchmark improvements on 31B-it + trtllm_mha + MTP + 1x B200 vs vLLM nightly (random 40 prompts x 1k/1k chat, max-concurrency=32): metric | SGLang (this PR) | vLLM nightly | Delta ------------------|------------------|--------------|---- outcome | OK | OK | same median TTFT | 673 ms | 901 ms | SGLang +25% median TPOT | 8.69 ms | 9.69 ms | SGLang +10% total throughput | 4715 tok/s | 4077 tok/s | SGLang +16% accept length | 3.13 | n/a | -- Same workload at conc=32 summarization (8k/1k x 40): median TPOT | 17.02 ms | 27.33 ms | SGLang +38% total throughput | 7475 tok/s | 6468 tok/s | SGLang +16% MMLU @ 500 questions on 31B-it: 0.680 vs vLLM 0.660 (within noise). Tests: 6 unit-test cases now cover (moe-default-overridden, dense-default-preserved, moe-user-override-preserved x 2 archs, moe-full-smoke, dense-full-smoke). Co-authored-by: Claude --- python/sglang/srt/server_args.py | 76 ++++++++----- test/srt/test_gemma4_swa_full_tokens_ratio.py | 106 +++++++++++++++--- 2 files changed, 140 insertions(+), 42 deletions(-) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 8878cfa36ccd..d4192f947744 100644 --- a/python/sglang/srt/server_args.py +++ b/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(): diff --git a/test/srt/test_gemma4_swa_full_tokens_ratio.py b/test/srt/test_gemma4_swa_full_tokens_ratio.py index 7a301cb557aa..70cff5be34b6 100644 --- a/test/srt/test_gemma4_swa_full_tokens_ratio.py +++ b/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"]))