diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 8303280d9..d9fb89b6d 100644 --- a/python/sglang/srt/environ.py +++ b/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 diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index bc2d7f183..ceae31dd5 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/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 ( - _is_npu - or _is_xpu - or supports_cuda_draft_extend_graph - or supports_hip_aiter_draft_extend_graph + if ( + self.draft_extend_attn_backend + and not envs.SGLANG_DISABLE_DRAFT_EXTEND_CUDA_GRAPH.get() + and ( + _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) diff --git a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py index de5d42685..f98298721 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py +++ b/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) diff --git a/test/registered/cp/test_deepseek_v4_flash_fp4_b200_cp.py b/test/registered/cp/test_deepseek_v4_flash_fp4_b200_cp.py index 5c0490221..815f36aa4 100644 --- a/test/registered/cp/test_deepseek_v4_flash_fp4_b200_cp.py +++ b/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", }