chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Log if cuda graph is used & extend cuda graph capture to cuda-graph-max-bs (#6201)

Browse Source 
Co-authored-by: SangBin Cho <rkooo567@gmail.com>
This commit is contained in:
Lianmin Zheng
2025-05-12 00:17:33 -07:00
committed by GitHub
parent 7d3a3d4510
commit fba8eccd7e
27 changed files with 293 additions and 121 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
python/sglang/srt/model_executor/cuda_graph_runner.py
View File

@@ -19,7 +19,7 @@ import bisect
import inspect
import os
from contextlib import contextmanager
from typing import TYPE_CHECKING, Callable
from typing import TYPE_CHECKING, Callable, Optional, Union
import torch
import tqdm
@@ -40,15 +40,12 @@ from sglang.srt.patch_torch import monkey_patch_torch_compile
from sglang.srt.utils import (
get_available_gpu_memory,
get_device_memory_capacity,
is_hip,
rank0_log,
)
if TYPE_CHECKING:
from sglang.srt.model_executor.model_runner import ModelRunner
_is_hip = is_hip()
def _to_torch(model: torch.nn.Module, reverse: bool, num_tokens: int):
for sub in model._modules.values():
@@ -137,7 +134,6 @@ def get_batch_sizes_to_capture(model_runner: ModelRunner):
)
gpu_mem = get_device_memory_capacity()
# Batch size of each rank will not become so large when DP is on
if gpu_mem is not None and gpu_mem > 96 * 1024:
capture_bs += list(range(160, 257, 8))
@@ -148,12 +144,15 @@ def get_batch_sizes_to_capture(model_runner: ModelRunner):
model_runner.req_to_token_pool.size
]
capture_bs = list(sorted(set(capture_bs)))
assert len(capture_bs) > 0 and capture_bs[0] > 0
capture_bs = [bs for bs in capture_bs if bs <= model_runner.req_to_token_pool.size]
if server_args.cuda_graph_max_bs:
capture_bs = [bs for bs in capture_bs if bs <= server_args.cuda_graph_max_bs]
if max(capture_bs) < server_args.cuda_graph_max_bs:
capture_bs += list(
range(max(capture_bs), server_args.cuda_graph_max_bs + 1, 16)
)
capture_bs = [bs for bs in capture_bs if bs <= model_runner.req_to_token_pool.size]
capture_bs = list(sorted(set(capture_bs)))
assert len(capture_bs) > 0 and capture_bs[0] > 0
compile_bs = (
[bs for bs in capture_bs if bs <= server_args.torch_compile_max_bs]
if server_args.enable_torch_compile