[Feat] Expose Encoder Mem Reserve As --encoder-mem-reserve CLI Flag (#339)

Author: zhaochenyang20

examples/run_qwen3_omni_server.py

@@ -69,6 +69,28 @@ def parse_args() -> argparse.Namespace:
             "If omitted, SGLang chooses automatically."
         ),
     )
+    parser.add_argument(
+        "--encoder-mem-reserve",
+        type=float,
+        default=None,
+        help=(
+            "GPU-memory fraction kept OUT of SGLang's static pool (model weights "
+            "+ KV cache) and left free for the co-located vision/audio encoder's "
+            "weights and activations on the thinker GPU.\n"
+            "Behavior across the four flag combinations of --mem-fraction-static "
+            "and --encoder-mem-reserve:\n"
+            "  (1) neither flag passed: SGLang auto-selects mem_fraction_static "
+            "and the default reserve 0.05 is subtracted;\n"
+            "  (2) only --encoder-mem-reserve X: SGLang auto-selects "
+            "mem_fraction_static and X is subtracted;\n"
+            "  (3) only --mem-fraction-static X: X is used verbatim and the "
+            "default reserve is ignored;\n"
+            "  (4) both flags: rejected at CLI as mutually exclusive.\n"
+            "Default 0.05 is tuned for single-request / short-video workloads; "
+            "raise to 0.15-0.20 for high-concurrency long-video or long-audio "
+            "workloads."
+        ),
+    )
 
     # Server
     parser.add_argument("--host", type=str, default="0.0.0.0")
@@ -83,14 +105,32 @@ def parse_args() -> argparse.Namespace:
     return parser.parse_args()
 
 
+def _check_mem_flag_mutex(
+    mem_fraction_static: float | None,
+    encoder_mem_reserve: float | None,
+) -> None:
+    """Reject passing both --mem-fraction-static and --encoder-mem-reserve."""
+    if mem_fraction_static is not None and encoder_mem_reserve is not None:
+        raise ValueError(
+            "--mem-fraction-static and --encoder-mem-reserve are mutually "
+            "exclusive: --mem-fraction-static pins the pool size directly "
+            "and the reserve only subtracts from SGLang's auto-selected "
+            "value. Pass only one."
+        )
+
+
 def main() -> None:
     args = parse_args()
 
+    _check_mem_flag_mutex(args.mem_fraction_static, args.encoder_mem_reserve)
+
     overrides = {}
     if args.thinker_max_seq_len is not None:
         overrides["thinker_max_seq_len"] = args.thinker_max_seq_len
     if args.cpu_offload_gb:
         overrides["cpu_offload_gb"] = args.cpu_offload_gb
+    if args.encoder_mem_reserve is not None:
+        overrides["encoder_mem_reserve"] = args.encoder_mem_reserve
 
     config = Qwen3OmniPipelineConfig(
         model_path=args.model_path,

sglang_omni/engines/ar/sglang_backend/server_args_builder.py

@@ -6,17 +6,6 @@
 
 from sglang.srt.server_args import ServerArgs
 
-# Note (Ratish, Chenyang):
-
-# SGLang's VLM auto-sizing applies a dynamic 0.95 * factor reserve
-# (roughly [0.8, 1.05]); Qwen3-Omni nests vision/audio configs under
-# `thinker_config` so SGLang's VLM path never triggers for us. 0.05
-# is a conservative linear lower-bound of that dynamic reserve; we
-# subtract it after auto-sizing when the thinker GPU also hosts encoder
-# stages. User-pinned mem_fraction_static bypasses this reserve.
-
-OMNI_ENCODER_MEM_FRACTION_STATIC_RESERVE = 0.05
-
 
 def build_sglang_server_args(
     model_path: str,
@@ -26,10 +15,9 @@ def build_sglang_server_args(
     max_prefill_tokens: int = 4096,
     max_running_requests: int = 16,
     mem_fraction_static: float | None = None,
-    auto_mem_fraction_static_reserve: float | None = None,
     **overrides: Any,
 ) -> ServerArgs:
-    """Build ServerArgs with shared defaults for all SGLang AR engines."""
+    """Build a SGLang ServerArgs with shared defaults for AR engines."""
     kwargs: dict[str, Any] = {
         "model_path": model_path,
         "trust_remote_code": True,
@@ -45,30 +33,36 @@ def build_sglang_server_args(
     if mem_fraction_static is not None:
         kwargs["mem_fraction_static"] = mem_fraction_static
     kwargs.update(overrides)
-    server_args = ServerArgs(**kwargs)
-    _apply_auto_mem_fraction_static_reserve(
-        server_args,
-        enabled=auto_mem_fraction_static_reserve is not None,
-        user_mem_fraction_static=mem_fraction_static,
-        reserve=auto_mem_fraction_static_reserve or 0.0,
-    )
-    return server_args
+    return ServerArgs(**kwargs)
 
 
-def _apply_auto_mem_fraction_static_reserve(
+def apply_encoder_mem_reserve(
     server_args: ServerArgs,
-    *,
-    enabled: bool,
-    user_mem_fraction_static: float | None,
-    reserve: float,
+    encoder_mem_reserve: float,
 ) -> None:
-    """Subtract a caller-requested reserve from SGLang's auto-selected value."""
-    if not enabled or user_mem_fraction_static is not None:
-        return
-    if reserve <= 0:
-        return
+    """Subtract encoder_mem_reserve from SGLang's auto-picked mem_fraction_static.
 
+    # Note (Chenyang):
+    Call this only when SGLang auto-selected mem_fraction_static —
+    i.e. the caller did NOT pin --mem-fraction-static. When the caller
+    pinned, that value is the whole budget and the reserve value is ignored.
+
+    Raises ValueError when the result would drop below 0.1 — below
+    that, SGLang's KV allocator fails deep in the scheduler with a
+    confusing traceback (empirically crashes ~0.08 on H200 for
+    Qwen3-Omni-30B), so surface it at build time instead.
+    """
+    if encoder_mem_reserve <= 0:
+        return
     current = server_args.mem_fraction_static
     if current is None:
         return
-    server_args.mem_fraction_static = round(max(0.01, current - reserve), 3)
+    new_value = current - encoder_mem_reserve
+    if new_value < 0.1:
+        raise ValueError(
+            f"auto mem_fraction_static {current:.3f} minus encoder_mem_reserve "
+            f"{encoder_mem_reserve:.3f} = {new_value:.3f} is below the safe "
+            f"floor 0.1; lower encoder_mem_reserve or pin "
+            f"--mem-fraction-static explicitly."
+        )
+    server_args.mem_fraction_static = round(new_value, 3)

sglang_omni/models/ming_omni/pipeline/stages.py

@@ -6,7 +6,7 @@
 from typing import Any
 
 from sglang_omni.engines.ar.sglang_backend.server_args_builder import (
-    OMNI_ENCODER_MEM_FRACTION_STATIC_RESERVE,
+    apply_encoder_mem_reserve,
     build_sglang_server_args,
 )
 from sglang_omni.engines.omni import create_sglang_ar_engine, create_single_pass_engine
@@ -338,9 +338,10 @@ def create_sglang_thinker_executor_from_config(
     server_args = build_sglang_server_args(
         local_path,
         context_length=thinker_max_seq_len,
-        auto_mem_fraction_static_reserve=OMNI_ENCODER_MEM_FRACTION_STATIC_RESERVE,
         **overrides,
     )
+    if "mem_fraction_static" not in overrides:
+        apply_encoder_mem_reserve(server_args, 0.05)
     pre_load_mem = (
         f", pre_load_avail_mem={pre_load_avail_mem:.2f} GB"
         if pre_load_avail_mem is not None

sglang_omni/models/qwen3_omni/config.py

@@ -113,27 +113,42 @@ def apply_server_args_overrides(
             and overrides["tp_size"] > 1
         ):
             raise NotImplementedError("Qwen3-Omni TP is not supported yet.")
-        remaining = _route_thinker_max_seq_len(self.stages, stage_name, overrides)
+        remaining = _route_thinker_executor_args(self.stages, stage_name, overrides)
         if remaining:
             super().apply_server_args_overrides(
                 stage_name=stage_name,
                 overrides=remaining,
             )
 
 
-def _route_thinker_max_seq_len(
+def _route_thinker_executor_args(
     stages: list[StageConfig],
     stage_name: str,
     overrides: dict[str, Any],
 ) -> dict[str, Any]:
+    """Pop thinker-factory kwargs onto the thinker stage; return the rest."""
     remaining = dict(overrides)
-    thinker_max_seq_len = remaining.pop("thinker_max_seq_len", None)
-    if thinker_max_seq_len is None or stage_name != THINKER_STAGE:
+    if stage_name != THINKER_STAGE:
         return remaining
-    for stage in stages:
-        if stage.name == THINKER_STAGE:
-            stage.executor.args["thinker_max_seq_len"] = int(thinker_max_seq_len)
-            break
+
+    casted: dict[str, Any] = {}
+
+    seq_len = remaining.pop("thinker_max_seq_len", None)
+    if seq_len is not None:
+        casted["thinker_max_seq_len"] = int(seq_len)
+
+    reserve = remaining.pop("encoder_mem_reserve", None)
+    if reserve is not None:
+        reserve = float(reserve)
+        if not 0.0 <= reserve < 1.0:
+            raise ValueError(f"encoder_mem_reserve must be in [0, 1), got {reserve}")
+        casted["encoder_mem_reserve"] = reserve
+
+    if casted:
+        for stage in stages:
+            if stage.name == THINKER_STAGE:
+                stage.executor.args.update(casted)
+                break
     return remaining
 
 
@@ -302,7 +317,7 @@ def apply_server_args_overrides(
             )
             if tp_size > 1:
                 raise NotImplementedError("Qwen3-Omni TP is not supported yet.")
-        remaining = _route_thinker_max_seq_len(self.stages, stage_name, overrides)
+        remaining = _route_thinker_executor_args(self.stages, stage_name, overrides)
         if remaining:
             super().apply_server_args_overrides(
                 stage_name=stage_name,

sglang_omni/models/qwen3_omni/pipeline/stages.py

@@ -10,7 +10,7 @@
 from transformers import AutoTokenizer
 
 from sglang_omni.engines.ar.sglang_backend.server_args_builder import (
-    OMNI_ENCODER_MEM_FRACTION_STATIC_RESERVE,
+    apply_encoder_mem_reserve,
     build_sglang_server_args,
 )
 from sglang_omni.engines.omni import (
@@ -354,21 +354,20 @@ def create_sglang_thinker_executor_from_config(
     *,
     gpu_id: int = 0,
     thinker_max_seq_len: int = 8192,
+    encoder_mem_reserve: float = 0.05,
     server_args_overrides: dict[str, Any] | None = None,
     speech_enabled: bool = False,
 ) -> EngineExecutor:
-    """Create a SGLang thinker executor from JSON-serializable config args.
-
-    This keeps pipeline config args plain dict types while still constructing
-    a typed ServerArgs object internally.
-    """
+    """Create a SGLang thinker executor from JSON-serializable config args."""
     pre_load_avail_mem = avail_gpu_mem(gpu_id)
+    overrides = server_args_overrides or {}
     server_args = build_sglang_server_args(
         model_path,
         context_length=thinker_max_seq_len,
-        auto_mem_fraction_static_reserve=OMNI_ENCODER_MEM_FRACTION_STATIC_RESERVE,
-        **(server_args_overrides or {}),
+        **overrides,
     )
+    if "mem_fraction_static" not in overrides:
+        apply_encoder_mem_reserve(server_args, encoder_mem_reserve)
     pre_load_mem = (
         f" pre_load_avail_mem={pre_load_avail_mem:.2f} GB"
         if pre_load_avail_mem is not None