diff --git a/docs/platforms/cpu_server.md b/docs/platforms/cpu_server.md index cf81a33bc..b954163e5 100644 --- a/docs/platforms/cpu_server.md +++ b/docs/platforms/cpu_server.md @@ -225,8 +225,7 @@ Notes: 3. For optimizing decoding with torch.compile, please add the flag `--enable-torch-compile`. To specify the maximum batch size when using `torch.compile`, set the flag `--torch-compile-max-bs`. For example, `--enable-torch-compile --torch-compile-max-bs 4` means using `torch.compile` - and setting the maximum batch size to 4. Currently the maximum applicable batch size - for optimizing with `torch.compile` is 16. + and setting the maximum batch size to 4. 4. A warmup step is automatically triggered when the service is started. The server is ready when you see the log `The server is fired up and ready to roll!`. diff --git a/python/sglang/srt/model_executor/cpu_graph_runner.py b/python/sglang/srt/model_executor/cpu_graph_runner.py index b98f531cd..8db80fc66 100644 --- a/python/sglang/srt/model_executor/cpu_graph_runner.py +++ b/python/sglang/srt/model_executor/cpu_graph_runner.py @@ -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 diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 319a1fd61..4f0f6c60d 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -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 diff --git a/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp b/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp index 67f68a569..88f4228a5 100644 --- a/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp +++ b/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp @@ -380,6 +380,9 @@ std::tuple image_preprocess_cpu( bool disable_grouping, at::ScalarType out_dtype); +// [NOTE] When registering kernels, we should accurately describe the in-place information. +// Taking fused_add_rmsnorm_cpu as an example, add `Tensor(a!)` modifier to all tensors that +// will be modified in-place to avoid incorrect fusing and execution order on graph mode. TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) { // activation m.def("silu_and_mul_cpu(Tensor input) -> Tensor"); diff --git a/test/srt/cpu/test_cpu_graph.py b/test/srt/cpu/test_cpu_graph.py index 1adc0e893..7e47c8663 100644 --- a/test/srt/cpu/test_cpu_graph.py +++ b/test/srt/cpu/test_cpu_graph.py @@ -31,9 +31,11 @@ class TestCPUGraph(CustomTestCase): "0.05", "--enable-torch-compile", "--torch-compile-max-bs", - "1", + "2", + "--cuda-graph-bs", + "2", ], - min_throughput=10, + min_throughput=7, ) def test_latency_torch_compile_cpu(self): return DEFAULT_MLA_MODEL_NAME_FOR_TEST @@ -58,8 +60,8 @@ class TestCPUGraph(CustomTestCase): "--trust-remote-code", "--disable-overlap-schedule", "--enable-torch-compile", - "--torch-compile-max-bs", - "1", + "--cuda-graph-bs", + "2", "--tp", f"{n_numa_node}", ],