[Spec] Add kill-switch env for draft-extend CUDA graph capture (#30944)

python/sglang/srt/environ.py

@@ -709,6 +709,11 @@ class Envs:
     # Saves the per-step draft forward, but the draft KV goes stale: an upshift
     # back to steps>0 starts from a cold draft state (low accept until it recovers).
     SGLANG_SPEC_SKIP_ZERO_STEP_DRAFT_EXTEND = EnvBool(False)
+    # Kill-switch for the draft-extend cuda graph. Draft extend then always runs
+    # eager. Escape hatch for setups where the capture's memory pool costs more
+    # than the graph saves (e.g. DeepEP MoE workspace captured at full dispatch
+    # capacity).
+    SGLANG_DISABLE_DRAFT_EXTEND_CUDA_GRAPH = EnvBool(False)
     # Use the split-KV (flash-decode) kernel for EAGLE target-verify on the
     # Triton backend (ROCm). Only active at speculative topk == 1; falls back to
     # extend_attention_fwd for unsupported cases or when set false (e.g. for

python/sglang/srt/speculative/eagle_worker_v2.py

@@ -486,11 +486,15 @@ class EagleDraftWorker(EagleDraftWorkerBase):
         ) and graph_supported_backend
         # Capture extend
         # TODO: support draft extend cuda graph for more attention backends
-        if self.draft_extend_attn_backend and (
+        if (
-            _is_npu
+            self.draft_extend_attn_backend
-            or _is_xpu
+            and not envs.SGLANG_DISABLE_DRAFT_EXTEND_CUDA_GRAPH.get()
-            or supports_cuda_draft_extend_graph
+            and (
-            or supports_hip_aiter_draft_extend_graph
+                _is_npu
+                or _is_xpu
+                or supports_cuda_draft_extend_graph
+                or supports_hip_aiter_draft_extend_graph
+            )
         ):
             tic = time.perf_counter()
             before_mem = get_available_gpu_memory(self.device, self.gpu_id)

python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py

@@ -237,6 +237,9 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase):
         if _is_cpu or check_cuda_graph_backend(Phase.DECODE, Backend.DISABLED):
             return
 
+        if envs.SGLANG_DISABLE_DRAFT_EXTEND_CUDA_GRAPH.get():
+            return
+
         if not _is_npu:
             self.cuda_graph_runner_for_draft_extend = (
                 MultiLayerEagleMultiStepDraftExtendCudaGraphRunner(self)

test/registered/cp/test_deepseek_v4_flash_fp4_b200_cp.py

@@ -30,6 +30,11 @@ DEEPEP_CONFIG = '{"normal_dispatch":{"num_sms":96},"normal_combine":{"num_sms":9
 
 _DEEPEP_ENV = {
     "SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "1024",
+    # The draft-extend graph pool costs ~4.5 GB here (DeepEP MoE workspace is
+    # captured at full dispatch capacity), which starves the eager prefill
+    # draft extend and OOMs. Run draft extend eager until the capture-time
+    # footprint is fixed.
+    "SGLANG_DISABLE_DRAFT_EXTEND_CUDA_GRAPH": "1",
 }