diff --git a/src/server/vllm_options.py b/src/server/vllm_options.py index 5bb35689..606ccbb4 100644 --- a/src/server/vllm_options.py +++ b/src/server/vllm_options.py @@ -76,6 +76,16 @@ def sampler_batch_limits(device_memory_bytes: int | None = None) -> dict[str, in } +def cudagraph_capture_sizes(max_num_seqs: int) -> list[int]: + """Batch sizes vLLM records a CUDA graph for: powers of two up to max_num_seqs. + + vLLM's default list has dozens of sizes, and capturing them (once to profile + memory, once for real) is a large part of a sampler's start-up. A batch runs + on the next captured size up. + """ + return [2**i for i in range(max_num_seqs.bit_length())] + + def detected_device_memory_bytes() -> int: try: import torch diff --git a/src/server/vllm_sampler.py b/src/server/vllm_sampler.py index ef627cbb..19fe43c9 100644 --- a/src/server/vllm_sampler.py +++ b/src/server/vllm_sampler.py @@ -26,7 +26,7 @@ from accel_timeslicer.time_slicer import TimeSlicerClient, time_slicer_client_from_env, time_slicing_enabled, workload_from_env from accel_timeslicer.workload import SAMPLER_CLAIM, WorkloadRef, local_workload_name from server.store import RequestStore, StateStore, get_state_store, get_store -from server.vllm_options import gpu_memory_utilization, sampler_batch_limits, split_stop, text_only_engine_kwargs +from server.vllm_options import cudagraph_capture_sizes, gpu_memory_utilization, sampler_batch_limits, split_stop, text_only_engine_kwargs tracer = trace.get_tracer("vllm.inference.worker") SHUTDOWN_SENTINEL = "SHUTDOWN_SENTINEL" @@ -43,17 +43,20 @@ def engine_kwargs_from_env(fft_enabled: bool) -> dict[str, Any]: model_name = os.getenv("BASE_MODEL") or os.getenv("VLLM_MODEL") if not model_name: raise ValueError("BASE_MODEL or VLLM_MODEL is required") + batch_limits = sampler_batch_limits() engine_kwargs = { "model": model_name, "enable_sleep_mode": fft_enabled, "enable_lora": not fft_enabled, "max_model_len": int(os.getenv("VLLM_MAX_MODEL_LEN", "8192")), - **sampler_batch_limits(), + **batch_limits, "gpu_memory_utilization": gpu_memory_utilization(), "enable_prefix_caching": False, "enforce_eager": os.getenv("VLLM_ENFORCE_EAGER", "0") == "1", **text_only_engine_kwargs(), } + if not engine_kwargs["enforce_eager"]: + engine_kwargs["compilation_config"] = {"cudagraph_capture_sizes": cudagraph_capture_sizes(batch_limits["max_num_seqs"])} architecture = os.getenv("VLLM_ARCHITECTURE_OVERRIDE") if architecture: engine_kwargs["hf_overrides"] = {"architectures": [architecture]} diff --git a/tests/test_vllm_options.py b/tests/test_vllm_options.py index d81f23dc..3629dfa4 100644 --- a/tests/test_vllm_options.py +++ b/tests/test_vllm_options.py @@ -2,7 +2,7 @@ import unittest from unittest.mock import patch -from server.vllm_options import gpu_memory_utilization, split_stop +from server.vllm_options import cudagraph_capture_sizes, gpu_memory_utilization, split_stop class SplitStopTest(unittest.TestCase): @@ -46,3 +46,9 @@ def test_placement_claim_does_not_cap_the_device(self) -> None: def test_without_a_budget_the_default_holds(self) -> None: with patch.dict(os.environ, {}, clear=True): self.assertEqual(gpu_memory_utilization(), 0.90) + + +class CudagraphCaptureSizesTest(unittest.TestCase): + def test_powers_of_two_up_to_max_num_seqs(self) -> None: + self.assertEqual(cudagraph_capture_sizes(256), [1, 2, 4, 8, 16, 32, 64, 128, 256]) + self.assertEqual(cudagraph_capture_sizes(96), [1, 2, 4, 8, 16, 32, 64])