diff --git a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py index dcd5f1435..34cc4043a 100644 --- a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py +++ b/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 diff --git a/test/registered/unit/model_executor/runner/test_prefill_cuda_graph_padding.py b/test/registered/unit/model_executor/runner/test_prefill_cuda_graph_padding.py new file mode 100644 index 000000000..8f8d80fa9 --- /dev/null +++ b/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()