diff --git a/python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py b/python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py index 2bd30e03d..475038f54 100644 --- a/python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py +++ b/python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py @@ -50,24 +50,6 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) -# Measured immediately before prefill graph construction, after model weights, -# KV cache, and the eager runner's static buffers have been allocated. Below -# this budget, compiling/capturing a multi-bucket prefill graph is likely to -# OOM or make no forward progress for large models. Keep an explicitly chosen -# backend untouched: an operator may intentionally trade KV capacity for it. -_MIN_AUTO_PREFILL_CUDA_GRAPH_FREE_MEMORY_GB = 4.0 - - -def should_skip_auto_prefill_cuda_graph_for_memory( - available_memory_gb: float, - cuda_graph_config_locked: set[tuple[str, str]], -) -> bool: - """Return whether an auto-selected prefill graph lacks capture headroom.""" - return ( - (Phase.PREFILL, "backend") not in cuda_graph_config_locked - and available_memory_gb < _MIN_AUTO_PREFILL_CUDA_GRAPH_FREE_MEMORY_GB - ) - class GraphCapture(msgspec.Struct, frozen=True, kw_only=True): runner: Optional[BaseRunner] @@ -396,20 +378,6 @@ def capture_prefill_graph( tic = time.perf_counter() before_mem = get_available_gpu_memory(model_runner.device, model_runner.gpu_id) - if should_skip_auto_prefill_cuda_graph_for_memory( - before_mem, - getattr(model_runner.server_args, "_cuda_graph_config_locked", set()), - ): - logger.warning( - "Disabling auto-selected prefill CUDA graph: only %.2f GiB is free " - "after model/KV/eager-buffer allocation; at least %.2f GiB is " - "required for capture. Set an explicit prefill CUDA graph backend " - "to override this safety gate.", - before_mem, - _MIN_AUTO_PREFILL_CUDA_GRAPH_FREE_MEMORY_GB, - ) - return result(eager_runner) - role = "draft" if model_runner.is_draft_worker else "target" capture_name = f"{role} prefill" logger.info( diff --git a/test/registered/unit/model_executor/model_runner_components/test_cuda_graph_setup.py b/test/registered/unit/model_executor/model_runner_components/test_cuda_graph_setup.py index 0632fe9c3..7e6a2ef1c 100644 --- a/test/registered/unit/model_executor/model_runner_components/test_cuda_graph_setup.py +++ b/test/registered/unit/model_executor/model_runner_components/test_cuda_graph_setup.py @@ -3,28 +3,15 @@ from types import SimpleNamespace import pytest -from sglang.srt.model_executor.cuda_graph_config import Phase from sglang.srt.model_executor.model_runner_components import cuda_graph_setup from sglang.srt.model_executor.model_runner_components.cuda_graph_setup import ( capture_decode_graph, - should_skip_auto_prefill_cuda_graph_for_memory, ) from sglang.test.ci.ci_register import register_cpu_ci register_cpu_ci(est_time=5, suite="base-a-test-cpu") -def test_auto_prefill_cuda_graph_memory_gate(): - assert should_skip_auto_prefill_cuda_graph_for_memory(3.99, set()) - assert not should_skip_auto_prefill_cuda_graph_for_memory(4.0, set()) - - -def test_explicit_prefill_backend_bypasses_memory_gate(): - assert not should_skip_auto_prefill_cuda_graph_for_memory( - 0.0, {(Phase.PREFILL, "backend")} - ) - - def test_model_runner_can_override_decode_graph_runner(monkeypatch): class CustomGraphRunner: def __init__(self, model_runner): diff --git a/test/registered/unit/model_executor/test_prefill_cuda_graph_runner.py b/test/registered/unit/model_executor/test_prefill_cuda_graph_runner.py index e52b08e0c..81baa61b8 100644 --- a/test/registered/unit/model_executor/test_prefill_cuda_graph_runner.py +++ b/test/registered/unit/model_executor/test_prefill_cuda_graph_runner.py @@ -61,6 +61,54 @@ class _FakeKVIndexKernel: class TestPrefillCudaGraphRunnerChunkedPrefix(CustomTestCase): + def test_low_free_memory_still_captures_prefill_graph(self): + eager_runner = object() + prefill_runner = object() + model_runner = SimpleNamespace( + device="cuda", + gpu_id=0, + is_draft_worker=False, + spec_algorithm=SimpleNamespace(is_eagle=lambda: False), + server_args=SimpleNamespace( + enable_lora=False, + cuda_graph_config=SimpleNamespace( + prefill=SimpleNamespace(bs=[1], backend=Backend.BREAKABLE) + ), + ), + model=SimpleNamespace(), + model_config=SimpleNamespace(context_len=8192, num_hidden_layers=1), + req_to_token_pool=SimpleNamespace(size=1), + ) + language_model = SimpleNamespace(layers=[object()]) + + with ( + patch.object(graph_setup, "check_cuda_graph_backend", return_value=False), + patch.object( + graph_setup, "resolve_language_model", return_value=language_model + ), + patch.object( + graph_setup, + "compute_attention_and_moe_layers", + return_value=([object()], [], [], [], []), + ), + patch.object( + graph_setup, + "get_available_gpu_memory", + side_effect=[3.99, 3.5], + ), + patch.object( + graph_setup, + "PrefillCudaGraphRunner", + return_value=prefill_runner, + ), + ): + capture = capture_prefill_graph( + model_runner=model_runner, + eager_runner=eager_runner, + ) + + self.assertIs(capture.runner, prefill_runner) + def test_eagle_target_tc_piecewise_skips_last_mode_capture(self): eager_runner = object() model_runner = SimpleNamespace(