From b425897366767d9f4a55e4bf67de7e6cb8d8b7c6 Mon Sep 17 00:00:00 2001
From: YC Yen-Ching Tseng <yctseng@amd.com>
Date: Tue, 1 Sep 2026 12:46:15 +0200
Subject: [PATCH] [AMD] Gate the aiter memory-reserve exemption behind an env
 var (#37242)

---
 .../docs/references/environment_variables.mdx |  5 ++++
 .../sglang/srt/arg_groups/attention_hook.py   | 25 +++++++++----------
 python/sglang/srt/environ.py                  |  8 ++++++
 3 files changed, 25 insertions(+), 13 deletions(-)

diff --git a/docs/docs/references/environment_variables.mdx b/docs/docs/references/environment_variables.mdx
index c0cb57e42..82e0ea4df 100644
--- a/docs/docs/references/environment_variables.mdx
+++ b/docs/docs/references/environment_variables.mdx
@@ -1802,6 +1802,11 @@ SGLang supports various environment variables that can be used to configure its
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Use AITER FP8 per-token quantization.</td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>false</code></td>
     </tr>
+    <tr>
+      <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_AITER_HONOR_EXPLICIT_MEM_FRACTION</code></td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Use an explicitly passed <code>--mem-fraction-static</code> as-is on AITER with context length above 8192, instead of scaling it by 0.85 to reserve non-static attention workspace. Skipping the reserve can OOM long-context serving; set it only when the scaled fraction is too small to hold the model weights.</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>false</code></td>
+    </tr>
     <tr>
       <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_USE_AITER_MOE_GU_ITLV</code></td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Select the AITER MoE gate/up tile layout: <code>true</code> interleaves, <code>false</code> uses the separated layout.</td>
diff --git a/python/sglang/srt/arg_groups/attention_hook.py b/python/sglang/srt/arg_groups/attention_hook.py
index 136b05956..282563416 100644
--- a/python/sglang/srt/arg_groups/attention_hook.py
+++ b/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,
                 )
diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py
index 8da36a531..70b2af61e 100644
--- a/python/sglang/srt/environ.py
+++ b/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).