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

Add support for more batch sizes in cpu_graph_runner (#13881)

Browse Source
This commit is contained in:
Cao E
2026-03-20 00:50:56 +08:00
committed by GitHub
parent 4c52b7fcc6
commit 274581fb77
5 changed files with 135 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

98
python/sglang/srt/model_executor/cpu_graph_runner.py
View File

@@ -17,6 +17,7 @@
from __future__ import annotations
import bisect
import logging
from contextlib import contextmanager
from typing import TYPE_CHECKING, Callable, Optional, Union
@@ -33,6 +34,7 @@ from sglang.srt.model_executor.forward_batch_info import (
ForwardBatch,
ForwardMode,
PPProxyTensors,
enable_num_token_non_padded,
)
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
from sglang.srt.utils import (
@@ -91,11 +93,15 @@ def set_torch_compile_config():
def get_batch_sizes_to_capture(model_runner: ModelRunner):
# torch compile speeds up decoding by reducing python overhead on CPU
server_args = model_runner.server_args
# cpu torch compile only speeds up decoding by
# reducing python overhead when bs is small
capture_bs = list(range(1, 17))
capture_bs = [bs for bs in capture_bs if bs <= server_args.torch_compile_max_bs]
# Note that we reuse server_args.cuda_graph_bs here.
# Users can customize the batch sizes supported by cpu_graph, such as:
# --cuda-graph-bs 1 2 4 8 16
capture_bs = server_args.cuda_graph_bs
assert (
max(capture_bs) <= server_args.torch_compile_max_bs
), f"{capture_bs=}, {server_args.torch_compile_max_bs=}"
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, f"{capture_bs=}"
@@ -520,6 +526,7 @@ class CPUGraphRunner:
# Batch sizes to capture
self.capture_bs = get_batch_sizes_to_capture(model_runner)
log_info_on_rank0(logger, f"Capture cpu graph bs {self.capture_bs}")
self.captured_forward_batches = {}
# Attention backend
self.max_bs = max(self.capture_bs)
self.max_num_token = self.max_bs * self.num_tokens_per_bs
@@ -564,7 +571,11 @@ class CPUGraphRunner:
)
def can_run(self, forward_batch: ForwardBatch):
is_bs_supported = forward_batch.batch_size in self.graphs
is_bs_supported = (
forward_batch.batch_size in self.graphs
if self.disable_padding
else forward_batch.batch_size <= self.max_bs
)
requested_capture_hidden_mode = max(
forward_batch.capture_hidden_mode,
@@ -669,8 +680,6 @@ class CPUGraphRunner:
num_token_non_padded=self.num_token_non_padded,
global_forward_mode=self.capture_forward_mode,
)
# Attention backend
self.model_runner.attn_backend.init_forward_metadata_capture_cpu_graph(
bs,
num_tokens,
@@ -705,6 +714,8 @@ class CPUGraphRunner:
for _ in range(2):
self.model_runner.tp_group.barrier()
out = run_once()
# Save the captured forward_batch
self.captured_forward_batches[bs] = forward_batch
return forward, out
def recapture_if_needed(self, forward_batch: ForwardBatch):
@@ -738,7 +749,53 @@ class CPUGraphRunner:
self.capture_hidden_mode = required_capture_hidden_mode
self.capture()
# TODO add padding support for CPUGraphRunner
def prepare_replay(
self,
forward_batch: ForwardBatch,
):
self.recapture_if_needed(forward_batch)
raw_bs = forward_batch.batch_size
if raw_bs in self.graphs:
self.model_runner.attn_backend.init_forward_metadata(forward_batch)
return forward_batch
raw_num_token = raw_bs * self.num_tokens_per_bs
index = bisect.bisect_left(self.capture_bs, raw_bs)
bs = self.capture_bs[index]
assert bs > raw_bs
self.raw_bs = raw_bs
self.raw_num_token = raw_num_token
self.bs = bs
captured_forward_batch = self.captured_forward_batches[bs]
assert captured_forward_batch is not None
captured_forward_batch.seq_lens.fill_(self.seq_len_fill_value)
captured_forward_batch.out_cache_loc.zero_()
captured_forward_batch.input_ids[:raw_num_token].copy_(forward_batch.input_ids)
captured_forward_batch.req_pool_indices[:raw_bs].copy_(
forward_batch.req_pool_indices
)
captured_forward_batch.seq_lens[:raw_bs].copy_(forward_batch.seq_lens)
captured_forward_batch.out_cache_loc[:raw_num_token].copy_(
forward_batch.out_cache_loc
)
captured_forward_batch.positions[:raw_num_token].copy_(forward_batch.positions)
if forward_batch.mrope_positions is not None:
self.mrope_positions[:, :raw_num_token].copy_(forward_batch.mrope_positions)
if self.is_encoder_decoder:
captured_forward_batch.encoder_lens[:raw_bs].copy_(
forward_batch.encoder_lens
)
if enable_num_token_non_padded(self.model_runner.server_args):
captured_forward_batch.num_token_non_padded.copy_(
forward_batch.num_token_non_padded
)
self.model_runner.attn_backend.init_forward_metadata(captured_forward_batch)
return captured_forward_batch
def replay(
self,
forward_batch: ForwardBatch,
@@ -748,14 +805,25 @@ class CPUGraphRunner:
assert (
pp_proxy_tensors is None
), "PPProxyTensors is not supported in CPUGraphRunner yet."
self.recapture_if_needed(forward_batch)
self.model_runner.attn_backend.init_forward_metadata(forward_batch)
output = self.graphs[forward_batch.batch_size](
forward_batch.input_ids,
forward_batch.positions,
forward_batch,
prepared_forward_batch = self.prepare_replay(forward_batch)
output = self.graphs[prepared_forward_batch.batch_size](
prepared_forward_batch.input_ids,
prepared_forward_batch.positions,
prepared_forward_batch,
)
if forward_batch.batch_size in self.graphs:
return output
assert isinstance(output, LogitsProcessorOutput)
return LogitsProcessorOutput(
next_token_logits=output.next_token_logits[: self.raw_num_token],
hidden_states=(
output.hidden_states[: self.raw_num_token]
if output.hidden_states is not None
else None
),
)
return output
def get_spec_info(self, num_tokens: int):
spec_info = None

46
python/sglang/srt/server_args.py
View File

@@ -1068,7 +1068,7 @@ class ServerArgs:
# 5. Pipeline parallelism
if self.pp_size > 1:
self.disable_piecewise_cuda_graph = True
# 6. Non-CUDA hardware (AMD, NPU, CPU, etc.)
# 6. Non-CUDA hardware (AMD, NPU, CPU, MPS, MUSA, XPU, etc.)
if is_hip() or is_npu() or is_cpu() or is_mps() or is_musa() or is_xpu():
self.disable_piecewise_cuda_graph = True
# 7. MoE A2A backend
@@ -1196,10 +1196,28 @@ class ServerArgs:
self.cuda_graph_max_bs = 160
# Set cuda graph batch sizes
if self.cuda_graph_bs is None:
self.cuda_graph_bs = self._generate_cuda_graph_batch_sizes()
if self.device != "cpu":
if self.cuda_graph_bs is None:
self.cuda_graph_bs = self._generate_cuda_graph_batch_sizes()
else:
self.cuda_graph_max_bs = max(self.cuda_graph_bs)
else:
self.cuda_graph_max_bs = max(self.cuda_graph_bs)
# Reuse cuda_graph_bs for cpu graph and use torch_compile_max_bs for cpu graph batch size limit,
# as cpu graph is based on torch.compile
if self.cuda_graph_bs is not None:
self.torch_compile_max_bs = max(self.cuda_graph_bs)
else:
# If cuda_graph_bs is not set, we will preferentially use torch_compile_max_bs
# to generate cuda_graph_bs
self.torch_compile_max_bs = (
self.torch_compile_max_bs or self.cuda_graph_max_bs
)
self.cuda_graph_bs = self._generate_cpu_graph_batch_sizes()
assert (
self.torch_compile_max_bs > 0
), "cuda_graph_bs should contain positive batch sizes"
self.cuda_graph_max_bs = self.torch_compile_max_bs
if self.piecewise_cuda_graph_max_tokens is None:
# Refer to pr #15927, by default we set the piecewise cuda graph max tokens to the chunked prefill size by default.
@@ -1321,6 +1339,26 @@ class ServerArgs:
return capture_bs
def _generate_cpu_graph_batch_sizes(self):
"""
Generate the list of batch sizes for CPU graph capture based on torch_compile_max_bs.
"""
if self.disable_cuda_graph_padding:
capture_bs = list(range(1, self.torch_compile_max_bs + 1))
else:
capture_bs = sorted(
set().union(
range(1, 17),
range(18, 31, 2),
range(32, 81, 4),
range(84, self.torch_compile_max_bs + 1, 8),
{self.torch_compile_max_bs},
)
)
capture_bs = [bs for bs in capture_bs if bs <= self.torch_compile_max_bs]
return capture_bs
def _generate_piecewise_cuda_graph_tokens(self):
"""
Generate the list of batch sizes for piecewise CUDA graph capture