[AMD] Gate the aiter memory-reserve exemption behind an env var (#37242)

python/sglang/srt/arg_groups/attention_hook.py

@@ -176,20 +176,19 @@ def handle_attention_backend_compatibility(server_args: Any):
     # AMD platforms backends
     if resolved_view(server_args).attention_backend == "aiter":
         if model_config.context_len > 8192:
-            # The 0.85 covers the extra non-static workspace aiter reserves for
-            # long contexts, but it is a heuristic for the auto-derived default
-            # only. Shrinking a value the user picked can push the static budget
-            # below the model-weight footprint on a nearly full GPU and break
-            # KV-cache allocation outright, so an explicit value is honored.
-            if (getattr(server_args, "_raw_input", None) or {}).get(
-                "mem_fraction_static"
-            ) is not None:
+            explicit_mem_fraction = (
+                getattr(server_args, "_raw_input", None) or {}
+            ).get("mem_fraction_static") is not None
+            if (
+                explicit_mem_fraction
+                and envs.SGLANG_AITER_HONOR_EXPLICIT_MEM_FRACTION.get()
+            ):
                 logger.warning(
-                    "attention_backend=aiter with context_len=%d (>8192) "
-                    "normally scales mem_fraction_static by 0.85, but "
-                    "mem_fraction_static=%.3f was set explicitly and will be "
-                    "used as-is. Ensure enough non-static memory is left for "
-                    "attention workspace and CUDA graphs.",
+                    "attention_backend=aiter with context_len=%d (>8192) normally "
+                    "scales mem_fraction_static by 0.85 to reserve non-static "
+                    "workspace, but SGLANG_AITER_HONOR_EXPLICIT_MEM_FRACTION is set, "
+                    "so mem_fraction_static=%.3f is used as-is. Ensure enough memory "
+                    "is left for attention workspace and CUDA graphs.",
                     model_config.context_len,
                     cfg.mem_fraction_static,
                 )

python/sglang/srt/environ.py

@@ -858,6 +858,14 @@ class Envs:
     SGLANG_ROCM_USE_MULTI_STREAM = EnvBool(False)
     SGLANG_HACK_FLASHMLA_BACKEND = EnvStr("tilelang")
     SGLANG_USE_AITER_FP8_PER_TOKEN = EnvBool(False)
+    # Above 8192 tokens of context, aiter's non-static workspace is large enough
+    # that mem_fraction_static is scaled by 0.85 to leave room for it. Set this to
+    # honor an explicitly passed --mem-fraction-static instead. Off by default:
+    # the reserve is load-bearing, and skipping it OOMs long-context aiter serving
+    # that fits comfortably with it (67.32 GiB request against 47.40 GiB free on a
+    # 288 GB MI355 in nightly-4-gpu-mi35x-minimax-m3). Worth setting only when the
+    # scaled fraction is itself too small to hold the model weights.
+    SGLANG_AITER_HONOR_EXPLICIT_MEM_FRACTION = EnvBool(False)
     # Route Kimi-K3-style h12 + fp8 MLA decode through aiter Triton Gluon when
     # import and Triton cga_layout prerequisites hold. Set to 0 to force the
     # zero-pad mla_decode_fwd fallback (benchmarking / emergency disable).