diff --git a/python/sglang/srt/compilation/compilation_config.py b/python/sglang/srt/compilation/compilation_config.py index 9da12bc91..d7687aaf1 100644 --- a/python/sglang/srt/compilation/compilation_config.py +++ b/python/sglang/srt/compilation/compilation_config.py @@ -18,6 +18,7 @@ class CompilationConfig: self.split_ops = [ "sglang.unified_attention_with_output", "sglang.gdn_with_output", + "sglang.inplace_all_reduce", ] def add_split_op(self, op: str): diff --git a/python/sglang/srt/compilation/cuda_piecewise_backend.py b/python/sglang/srt/compilation/cuda_piecewise_backend.py index dd8b358b7..1955fd3d5 100644 --- a/python/sglang/srt/compilation/cuda_piecewise_backend.py +++ b/python/sglang/srt/compilation/cuda_piecewise_backend.py @@ -11,6 +11,10 @@ import torch.fx as fx from sglang.srt.compilation.compilation_config import CompilationConfig from sglang.srt.compilation.compilation_counter import compilation_counter +from sglang.srt.compilation.piecewise_context_manager import ( + get_pcg_capture_stream, + is_in_pcg_torch_compile, +) from sglang.srt.compilation.weak_ref_tensor import weak_ref_tensors logger = logging.getLogger(__name__) @@ -105,6 +109,10 @@ class CUDAPiecewiseBackend: self.first_run_finished = True self.check_for_ending_compilation() return self.compiled_graph_for_general_shape(*args) + + if len(self.sym_shape_indices) == 0: + return self.compiled_graph_for_general_shape(*args) + runtime_shape = args[self.sym_shape_indices[0]] if runtime_shape not in self.concrete_size_entries: # we don't need to do anything for this shape @@ -137,6 +145,8 @@ class CUDAPiecewiseBackend: # skip_cuda_graphs = get_forward_context().skip_cuda_graphs # if not entry.use_cudagraph or skip_cuda_graphs: # return entry.runnable(*args) + if is_in_pcg_torch_compile(): + return entry.runnable(*args) if entry.cudagraph is None: if entry.num_finished_warmup < 1: # noqa @@ -160,9 +170,10 @@ class CUDAPiecewiseBackend: # and disable gc for the rest of the graphs. stack.enter_context(patch("gc.collect", lambda: None)) stack.enter_context(patch("torch.cuda.empty_cache", lambda: None)) - # mind-exploding: carefully manage the reference and memory. - with torch.cuda.graph(cudagraph, pool=self.graph_pool): + stream = get_pcg_capture_stream() + assert stream is not None, "PCG capture stream is not set" + with torch.cuda.graph(cudagraph, pool=self.graph_pool, stream=stream): # `output` is managed by pytorch's cudagraph pool output = entry.runnable(*args) if self.is_last_graph: @@ -194,6 +205,5 @@ class CUDAPiecewiseBackend: "Input addresses for cudagraphs are different during replay." f" Expected {entry.input_addresses}, got {new_input_addresses}" ) - entry.cudagraph.replay() return entry.output diff --git a/python/sglang/srt/compilation/piecewise_context_manager.py b/python/sglang/srt/compilation/piecewise_context_manager.py index bafb5645d..d8f6f5cbe 100644 --- a/python/sglang/srt/compilation/piecewise_context_manager.py +++ b/python/sglang/srt/compilation/piecewise_context_manager.py @@ -2,15 +2,35 @@ from contextlib import contextmanager from dataclasses import dataclass from typing import Any, List, Optional +import torch + from sglang.srt.model_executor.forward_batch_info import ForwardBatch _in_piecewise_cuda_graph = False +_in_pcg_torch_compile = False +_pcg_capture_stream = None def is_in_piecewise_cuda_graph(): return _in_piecewise_cuda_graph +def is_in_pcg_torch_compile(): + return _in_pcg_torch_compile + + +def get_pcg_capture_stream(): + return _pcg_capture_stream + + +@contextmanager +def enable_piecewise_cuda_graph_compile(): + global _in_pcg_torch_compile + _in_pcg_torch_compile = True + yield + _in_pcg_torch_compile = False + + @contextmanager def enable_piecewise_cuda_graph(): global _in_piecewise_cuda_graph @@ -21,6 +41,14 @@ def enable_piecewise_cuda_graph(): _in_piecewise_cuda_graph = False +@contextmanager +def set_pcg_capture_stream(stream: torch.cuda.Stream): + global _pcg_capture_stream + _pcg_capture_stream = stream + yield + _pcg_capture_stream = None + + @dataclass class ForwardContext: def __init__(self): diff --git a/python/sglang/srt/managers/mm_utils.py b/python/sglang/srt/managers/mm_utils.py index fea6bee0c..1d8e48fbd 100644 --- a/python/sglang/srt/managers/mm_utils.py +++ b/python/sglang/srt/managers/mm_utils.py @@ -13,7 +13,6 @@ import numpy as np import torch from torch import nn -from sglang.srt.distributed.parallel_state import get_tp_group from sglang.srt.environ import envs from sglang.srt.layers.multimodal import gpu_tensor_hash from sglang.srt.managers.schedule_batch import ( @@ -1043,69 +1042,52 @@ def general_mm_embed_routine( Returns: Hidden states from language model forward pass """ - # Lazy import to allow some monkey patch of piecewise_cuda_graph_runner - from sglang.srt.model_executor.piecewise_cuda_graph_runner import ( - use_original_ca_comm, - ) - - tp_group = get_tp_group() - - with use_original_ca_comm(tp_group): - # We disable custom allreduce in piecewise cuda graph. - # However, because we only capture the language model part, the multimodal can still use custom allreduce. - assert hasattr(language_model, "get_input_embeddings") - embed_tokens = language_model.get_input_embeddings() + assert hasattr(language_model, "get_input_embeddings") + embed_tokens = language_model.get_input_embeddings() + if not hasattr(language_model, "pp_group") or language_model.pp_group.is_first_rank: if ( - not hasattr(language_model, "pp_group") - or language_model.pp_group.is_first_rank + not forward_batch.forward_mode.is_decode() + and not forward_batch.forward_mode.is_target_verify() + and forward_batch.contains_mm_inputs() ): - if ( - not forward_batch.forward_mode.is_decode() - and not forward_batch.forward_mode.is_target_verify() - and forward_batch.contains_mm_inputs() - ): - mm_inputs_list = [ - mm_input - for mm_input in forward_batch.mm_inputs - if mm_input is not None - ] - extend_prefix_lens = [ - prefix_len - for i, prefix_len in enumerate(forward_batch.extend_prefix_lens_cpu) - if forward_batch.mm_inputs[i] is not None - ] - extend_seq_lens = [ - seq_len - for i, seq_len in enumerate(forward_batch.extend_seq_lens_cpu) - if forward_batch.mm_inputs[i] is not None - ] - input_embeds, other_info = embed_mm_inputs( - mm_inputs_list=mm_inputs_list, - extend_prefix_lens=extend_prefix_lens, - extend_seq_lens=extend_seq_lens, - input_ids=input_ids, - multimodal_model=multimodal_model, - input_embedding=embed_tokens, - data_embedding_func_mapping=data_embedding_funcs, - placeholder_tokens=placeholder_tokens, - use_deepstack=use_deepstack, - ) - # add for qwen3_vl deepstack - if use_deepstack: - kwargs["input_deepstack_embeds"] = other_info[ - "input_deepstack_embeds" - ] - # once used, mm_inputs is useless, considering chunked-prefill is disabled for multimodal models - # just being defensive here - forward_batch.mm_inputs = None - else: - input_embeds = embed_tokens(input_ids) - # Copy to pre-allocated buffer if available (for CUDA graph address stability) - if forward_batch.input_embeds is not None: - forward_batch.input_embeds.copy_(input_embeds) - input_embeds = forward_batch.input_embeds + mm_inputs_list = [ + mm_input for mm_input in forward_batch.mm_inputs if mm_input is not None + ] + extend_prefix_lens = [ + prefix_len + for i, prefix_len in enumerate(forward_batch.extend_prefix_lens_cpu) + if forward_batch.mm_inputs[i] is not None + ] + extend_seq_lens = [ + seq_len + for i, seq_len in enumerate(forward_batch.extend_seq_lens_cpu) + if forward_batch.mm_inputs[i] is not None + ] + input_embeds, other_info = embed_mm_inputs( + mm_inputs_list=mm_inputs_list, + extend_prefix_lens=extend_prefix_lens, + extend_seq_lens=extend_seq_lens, + input_ids=input_ids, + multimodal_model=multimodal_model, + input_embedding=embed_tokens, + data_embedding_func_mapping=data_embedding_funcs, + placeholder_tokens=placeholder_tokens, + use_deepstack=use_deepstack, + ) + # add for qwen3_vl deepstack + if use_deepstack: + kwargs["input_deepstack_embeds"] = other_info["input_deepstack_embeds"] + # once used, mm_inputs is useless, considering chunked-prefill is disabled for multimodal models + # just being defensive here + forward_batch.mm_inputs = None else: - input_embeds = None + input_embeds = embed_tokens(input_ids) + # Copy to pre-allocated buffer if available (for CUDA graph address stability) + if forward_batch.input_embeds is not None: + forward_batch.input_embeds.copy_(input_embeds) + input_embeds = forward_batch.input_embeds + else: + input_embeds = None hidden_states = language_model( input_ids=None, diff --git a/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py b/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py index 94a16228b..2dcd59ceb 100644 --- a/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py @@ -29,13 +29,16 @@ from sglang.srt.compilation.compilation_config import CompilationConfig from sglang.srt.compilation.compile import install_torch_compiled, set_compiled from sglang.srt.compilation.piecewise_context_manager import ( enable_piecewise_cuda_graph, + enable_piecewise_cuda_graph_compile, set_forward_context, + set_pcg_capture_stream, ) from sglang.srt.custom_op import CustomOp from sglang.srt.distributed import get_tensor_model_parallel_rank from sglang.srt.distributed.device_communicators.pynccl_allocator import ( set_graph_pool_id, ) +from sglang.srt.distributed.parallel_state import graph_capture from sglang.srt.layers.dp_attention import ( DpPaddingMode, get_attention_tp_rank, @@ -60,50 +63,6 @@ if TYPE_CHECKING: from sglang.srt.model_executor.model_runner import ModelRunner -@contextmanager -def disable_ca_comm(tp_group): - """ - Context manager to temporarily disable custom allreduce communication. - - This is used during Piecewise CUDA graph capture to avoid custom allreduce operations - that may not be compatible with graph capture. - - TODO(yuwei): Fix this - """ - if tp_group.ca_comm is None: - yield - return - - original_disabled = tp_group.ca_comm.disabled - tp_group.ca_comm.original_disabled = original_disabled - try: - tp_group.ca_comm.disabled = True - yield - finally: - tp_group.ca_comm.disabled = original_disabled - - -@contextmanager -def use_original_ca_comm(tp_group): - """ - For the module not in piecewise cuda graph capture, use the original custom allreduce communication. - This is a no-op if not using piecewise cuda graph because .disabled == .original_disabled - - TODO(Byron): remove this once custom allreduce is enabled in piecewise cuda graph - """ - if tp_group.ca_comm is None: - yield - return - - current_disabled = tp_group.ca_comm.disabled - original_disabled = tp_group.ca_comm.original_disabled - try: - tp_group.ca_comm.disabled = original_disabled - yield - finally: - tp_group.ca_comm.disabled = current_disabled - - @contextmanager def freeze_gc(enable_cudagraph_gc: bool): """ @@ -267,9 +226,25 @@ class PiecewiseCudaGraphRunner: compile_config=self.compile_config, graph_pool=get_global_graph_memory_pool(), ) - with freeze_gc(self.model_runner.server_args.enable_cudagraph_gc): - with set_compiled(True): - self.warmup_torch_compile() + + with set_compiled(True), enable_piecewise_cuda_graph_compile(): + compile_range = ( + tqdm.tqdm(list(reversed(self.capture_num_tokens))) + if get_tensor_model_parallel_rank() == 0 + else reversed(self.capture_num_tokens) + ) + for _, num_tokens in enumerate(compile_range): + if get_tensor_model_parallel_rank() == 0: + compile_range.set_description( + f"Compiling num tokens ({num_tokens=})" + ) + self.warmup_torch_compile(num_tokens=num_tokens) + + set_global_graph_memory_pool(self.device_module.graph_pool_handle()) + set_graph_pool_id(get_global_graph_memory_pool()) + + self.device_module.synchronize() + self.model_runner.tp_group.barrier() # Capture try: self.capture() @@ -280,21 +255,20 @@ class PiecewiseCudaGraphRunner: self.raw_num_tokens = 0 - def warmup_torch_compile(self): + def warmup_torch_compile(self, num_tokens: int): """Warmup the model with a simple forward pass before CUDA graph capture.""" - num_tokens = 2 input_ids = self.input_ids[:num_tokens] input_embeds = self.input_embeds[:num_tokens] if self.is_multimodal else None + positions = self.positions[:num_tokens] + mrope_positions = ( + self.mrope_positions[:, :num_tokens] if self.is_multimodal else None + ) out_cache_loc = self.out_cache_loc[:num_tokens] out_cache_loc_swa = ( self.out_cache_loc_swa[:num_tokens] if self.out_cache_loc_swa is not None else None ) - positions = self.positions[:num_tokens] - mrope_positions = ( - self.mrope_positions[:, :num_tokens] if self.is_multimodal else None - ) with torch.device(self.device): forward_batch = ForwardBatch( forward_mode=ForwardMode.EXTEND, @@ -337,9 +311,12 @@ class PiecewiseCudaGraphRunner: # Attention backend self.model_runner.attn_backend.init_forward_metadata(forward_batch) + forward_batch.dp_local_start_pos = forward_batch.dp_local_num_tokens = None + set_dp_buffer_len(None, num_tokens, forward_batch.dp_padding_mode.is_max_len()) + set_is_extend_in_batch(False) with set_forward_context( forward_batch, self.attention_layers, self.quant_config, self.moe_layers - ), disable_ca_comm(self.model_runner.tp_group): + ): _ = self.model_runner.model.forward( forward_batch.input_ids, forward_batch.positions, @@ -351,7 +328,6 @@ class PiecewiseCudaGraphRunner: def can_run(self, forward_batch: ForwardBatch): num_tokens = len(forward_batch.input_ids) - # TODO(yuwei): support return input_ids' logprob if forward_batch.return_logprob: for start_len, seq_len in zip( forward_batch.extend_logprob_start_lens_cpu, @@ -369,31 +345,33 @@ class PiecewiseCudaGraphRunner: # can reuse the memory pool allocated for the large shapes. with freeze_gc( self.model_runner.server_args.enable_cudagraph_gc - ), disable_ca_comm(self.model_runner.tp_group): - avail_mem = get_available_gpu_memory( - self.model_runner.device, - self.model_runner.gpu_id, - empty_cache=False, - ) - # Reverse the order to enable better memory sharing across cuda graphs. - capture_range = ( - tqdm.tqdm(list(reversed(self.capture_num_tokens))) - if get_tensor_model_parallel_rank() == 0 - else reversed(self.capture_num_tokens) - ) - for i, num_tokens in enumerate(capture_range): - if get_tensor_model_parallel_rank() == 0: - avail_mem = get_available_gpu_memory( - self.model_runner.device, - self.model_runner.gpu_id, - empty_cache=False, - ) - capture_range.set_description( - f"Capturing num tokens ({num_tokens=} {avail_mem=:.2f} GB)" - ) + ), graph_capture() as graph_capture_context: + stream = graph_capture_context.stream + with set_pcg_capture_stream(stream): + avail_mem = get_available_gpu_memory( + self.model_runner.device, + self.model_runner.gpu_id, + empty_cache=False, + ) + # Reverse the order to enable better memory sharing across cuda graphs. + capture_range = ( + tqdm.tqdm(list(reversed(self.capture_num_tokens))) + if get_tensor_model_parallel_rank() == 0 + else reversed(self.capture_num_tokens) + ) + for i, num_tokens in enumerate(capture_range): + if get_tensor_model_parallel_rank() == 0: + avail_mem = get_available_gpu_memory( + self.model_runner.device, + self.model_runner.gpu_id, + empty_cache=False, + ) + capture_range.set_description( + f"Capturing num tokens ({num_tokens=} {avail_mem=:.2f} GB)" + ) - with set_compiled(True): - self.capture_one_batch_size(num_tokens) + with set_compiled(True): + self.capture_one_batch_size(num_tokens) def capture_one_batch_size(self, num_tokens: int): bs = 1 @@ -493,7 +471,9 @@ class PiecewiseCudaGraphRunner: ) return - for _ in range(3): + # run twice for warmup at the first time and cuda graph capture at the second time + # detail lies in sglang/python/sglang/srt/compilation/cuda_piecewise_backend.py + for _ in range(2): self.device_module.synchronize() self.model_runner.tp_group.barrier() run_once() @@ -603,7 +583,7 @@ class PiecewiseCudaGraphRunner: forward_batch: ForwardBatch, **kwargs, ) -> Union[LogitsProcessorOutput, PPProxyTensors, EmbeddingPoolerOutput]: - with enable_piecewise_cuda_graph(), disable_ca_comm(self.model_runner.tp_group): + with enable_piecewise_cuda_graph(): self.model_runner.attn_backend.init_forward_metadata(forward_batch) static_forward_batch = self.replay_prepare(forward_batch, **kwargs) # Replay