perf(sampler): capture CUDA graphs for power-of-two batch sizes only

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

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]}

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])