Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions python/sglang/srt/server_args.py
Original file line number Diff line number Diff line change
Expand Up @@ -1299,6 +1299,9 @@ def _generate_cuda_graph_batch_sizes(self):

capture_bs = [bs for bs in capture_bs if bs <= self.cuda_graph_max_bs]

if self.cuda_graph_max_bs not in capture_bs:
capture_bs.append(self.cuda_graph_max_bs)
Comment thread
weireweire marked this conversation as resolved.

return capture_bs

def _generate_piecewise_cuda_graph_tokens(self):
Expand Down
10 changes: 10 additions & 0 deletions test/registered/core/test_server_args.py
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,16 @@ def test_pd_decode_defaults_to_round_robin(self):
self.assertEqual(server_args.load_balance_method, "round_robin")


class TestCudaGraphBatchSizes(unittest.TestCase):
def test_generate_cuda_graph_batch_sizes_includes_max_bs(self):
server_args = ServerArgs(model_path="dummy", cuda_graph_max_bs=500)

capture_bs = server_args._generate_cuda_graph_batch_sizes()

self.assertIn(500, capture_bs)
self.assertEqual(capture_bs[-1], 500)


class TestPortArgs(unittest.TestCase):
@patch("sglang.srt.server_args.get_free_port")
@patch("sglang.srt.server_args.tempfile.NamedTemporaryFile")
Expand Down
Loading