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

Custom All Reduce for Piecewise Cuda Graph (#15356)

Browse Source 
Signed-off-by: Oasis-Git <ayw.sirius19@gmail.com>
This commit is contained in:
Yuwei An
2025-12-25 07:55:15 -08:00
committed by GitHub
parent b6702d72cf
commit 5c243ba588
5 changed files with 146 additions and 145 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
python/sglang/srt/compilation/compilation_config.py
View File

@@ -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):

16
python/sglang/srt/compilation/cuda_piecewise_backend.py
View File

@@ -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

28
python/sglang/srt/compilation/piecewise_context_manager.py
View File

@@ -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):

104
python/sglang/srt/managers/mm_utils.py
View File

@@ -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,

142
python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py
View File

@@ -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