perf: avoid excessive prefill CUDA graph padding (#31487)

python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py

@@ -108,6 +108,11 @@ from sglang.srt.utils import (
 warnings.filterwarnings("ignore", message=".*lru_cache.*", module="torch._dynamo")
 logger = logging.getLogger(__name__)
 
+# A replay executes every padded token in its capture bucket. Sparse bucket
+# lists can otherwise turn the lower launch overhead into substantially more
+# model work than an exact-shape eager forward.
+_MAX_PREFILL_CUDA_GRAPH_PADDING_FACTOR = 2
+
 _is_hip = is_hip()
 _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
 
@@ -662,10 +667,12 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
                     return False
         if num_tokens > self.max_num_tokens:
             return False
-        # No backend-level shape check here: load_batch bucket-pads
-        # num_tokens up to the nearest captured shape, so eligibility is
-        # bounded by num_tokens <= self.max_num_tokens (already
-        # checked above), not by exact shape membership.
+        padded_num_tokens = self._pad_to_bucket(num_tokens, self.capture_num_tokens)
+        if padded_num_tokens > num_tokens * _MAX_PREFILL_CUDA_GRAPH_PADDING_FACTOR:
+            return False
+        # No exact-shape check here: load_batch bucket-pads to the nearest
+        # captured shape. The factor above only rejects replays whose padded
+        # model work is disproportionate to the useful token count.
         #
         # Multi-req replay is supported by BCG via the layer_model.forward
         # monkey-patch in replay(): the captured bs=1 graph runs the

test/registered/unit/model_executor/runner/test_prefill_cuda_graph_padding.py

@@ -0,0 +1,51 @@
+import unittest
+from types import SimpleNamespace
+
+from sglang.srt.model_executor.forward_batch_info import (
+    CaptureHiddenMode,
+    ForwardMode,
+)
+from sglang.srt.model_executor.runner.prefill_cuda_graph_runner import (
+    PrefillCudaGraphRunner,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=1, suite="base-a-test-cpu")
+
+
+class TestPrefillCudaGraphPadding(CustomTestCase):
+    def _make_runner(self):
+        runner = PrefillCudaGraphRunner.__new__(PrefillCudaGraphRunner)
+        runner._is_full_backend = False
+        runner.capture_hidden_mode = CaptureHiddenMode.NULL
+        runner.capture_num_tokens = [4, 16]
+        runner.max_num_tokens = 16
+        return runner
+
+    def _make_forward_batch(self, num_tokens):
+        return SimpleNamespace(
+            batch_size=1,
+            input_embeds=None,
+            replace_embeds=None,
+            mm_inputs=None,
+            forward_mode=ForwardMode.EXTEND,
+            capture_hidden_mode=CaptureHiddenMode.NULL,
+            global_num_tokens_cpu=None,
+            return_logprob=False,
+            input_ids=list(range(num_tokens)),
+        )
+
+    def test_rejects_more_than_two_x_token_padding(self):
+        runner = self._make_runner()
+
+        self.assertFalse(runner.can_run_graph(self._make_forward_batch(5)))
+
+    def test_accepts_two_x_token_padding(self):
+        runner = self._make_runner()
+
+        self.assertTrue(runner.can_run_graph(self._make_forward_batch(8)))
+
+
+if __name__ == "__main__":
+    unittest.main()