1443 lines
56 KiB
Python
1443 lines
56 KiB
Python
# Copyright 2023-2024 SGLang Team
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
# ==============================================================================
|
|
"""Run the model with cuda graph and torch.compile."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import bisect
|
|
import contextlib
|
|
import gc
|
|
import inspect
|
|
import logging
|
|
import os
|
|
from contextlib import contextmanager
|
|
from dataclasses import dataclass
|
|
from functools import partial
|
|
from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Tuple, Union
|
|
|
|
import torch
|
|
import tqdm
|
|
from torch.profiler import ProfilerActivity, profile
|
|
|
|
from sglang.srt.batch_overlap.two_batch_overlap import TboCudaGraphRunnerPlugin
|
|
from sglang.srt.constants import GPU_MEMORY_TYPE_CUDA_GRAPH
|
|
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 (
|
|
GroupCoordinator,
|
|
graph_capture,
|
|
set_pdmux_status,
|
|
)
|
|
from sglang.srt.dllm.config import DllmConfig
|
|
from sglang.srt.environ import envs
|
|
from sglang.srt.layers.attention.dsa.utils import is_dsa_enable_prefill_cp
|
|
from sglang.srt.layers.dp_attention import (
|
|
DpPaddingMode,
|
|
get_attention_cp_size,
|
|
get_attention_tp_rank,
|
|
get_attention_tp_size,
|
|
set_dp_buffer_len,
|
|
set_is_extend_in_batch,
|
|
)
|
|
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
|
|
from sglang.srt.layers.moe.token_dispatcher.deepep import DeepEPBuffer
|
|
from sglang.srt.layers.moe.utils import get_deepep_mode, get_moe_a2a_backend
|
|
from sglang.srt.layers.utils import MultiPlatformOp
|
|
from sglang.srt.layers.utils.cp_utils import is_mla_prefill_cp_enabled
|
|
from sglang.srt.model_executor.forward_batch_info import (
|
|
CaptureHiddenMode,
|
|
ForwardBatch,
|
|
ForwardMode,
|
|
NgramEmbeddingInfo,
|
|
PPProxyTensors,
|
|
compute_local_num_token_non_padded,
|
|
enable_num_token_non_padded,
|
|
)
|
|
from sglang.srt.model_executor.forward_context import ForwardContext, forward_context
|
|
from sglang.srt.model_executor.input_buffers import ForwardInputBuffers
|
|
from sglang.srt.multiplex.pdmux_context import get_current_stream_idx, get_stream_groups
|
|
from sglang.srt.utils import (
|
|
empty_context,
|
|
get_available_gpu_memory,
|
|
get_bool_env_var,
|
|
is_hip,
|
|
log_info_on_rank0,
|
|
require_attn_tp_gather,
|
|
require_gathered_buffer,
|
|
require_mlp_sync,
|
|
require_mlp_tp_gather,
|
|
)
|
|
from sglang.srt.utils.patch_torch import monkey_patch_torch_compile
|
|
from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter
|
|
|
|
try:
|
|
from kt_kernel import KTMoEWrapper
|
|
|
|
KTRANSFORMERS_AVAILABLE = True
|
|
except ImportError:
|
|
KTRANSFORMERS_AVAILABLE = False
|
|
|
|
_is_hip = is_hip()
|
|
|
|
if not _is_hip:
|
|
from sglang.srt.model_executor.breakable_cuda_graph.breakable_cuda_graph import (
|
|
BreakableCUDAGraph,
|
|
BreakableCUDAGraphCapture,
|
|
eager_on_graph,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
if TYPE_CHECKING:
|
|
from sglang.srt.model_executor.model_runner import ModelRunner
|
|
|
|
_has_foreach_copy = hasattr(torch, "_foreach_copy_")
|
|
|
|
|
|
def _grouped_foreach_copy_(dsts: List[torch.Tensor], srcs: List[torch.Tensor]) -> None:
|
|
"""Call torch._foreach_copy_ grouped by (dst_dtype, src_dtype) pairs."""
|
|
|
|
def foreach_copy(dsts: List[torch.Tensor], srcs: List[torch.Tensor]) -> None:
|
|
if _has_foreach_copy:
|
|
torch._foreach_copy_(dsts, srcs)
|
|
else:
|
|
for dst, src in zip(dsts, srcs):
|
|
dst.copy_(src)
|
|
|
|
groups: Dict[Tuple[torch.dtype, torch.dtype], Tuple[List, List]] = {}
|
|
for dst, src in zip(dsts, srcs):
|
|
key = (dst.dtype, src.dtype)
|
|
if key not in groups:
|
|
groups[key] = ([], [])
|
|
groups[key][0].append(dst)
|
|
groups[key][1].append(src)
|
|
for group_dsts, group_srcs in groups.values():
|
|
foreach_copy(group_dsts, group_srcs)
|
|
|
|
|
|
@dataclass
|
|
class DecodeInputBuffers(ForwardInputBuffers):
|
|
|
|
input_ids: torch.Tensor
|
|
input_embeds: torch.Tensor
|
|
req_pool_indices: torch.Tensor
|
|
seq_lens: torch.Tensor
|
|
seq_lens_cpu: torch.Tensor
|
|
out_cache_loc: torch.Tensor
|
|
positions: torch.Tensor
|
|
mrope_positions: torch.Tensor
|
|
num_token_non_padded: torch.Tensor
|
|
custom_mask: torch.Tensor
|
|
next_token_logits_buffer: torch.Tensor
|
|
mamba_track_indices: Optional[torch.Tensor]
|
|
mamba_track_mask: Optional[torch.Tensor]
|
|
global_num_tokens_gpu: torch.Tensor
|
|
global_num_tokens_for_logprob_gpu: torch.Tensor
|
|
encoder_lens: Optional[torch.Tensor]
|
|
pp_proxy_tensors: Optional[Dict[str, torch.Tensor]]
|
|
ngram_embedding_info: Optional["NgramEmbeddingInfo"]
|
|
|
|
@classmethod
|
|
def create(
|
|
cls,
|
|
*,
|
|
device: torch.device,
|
|
max_bs: int,
|
|
max_num_token: int,
|
|
hidden_size: int,
|
|
vocab_size: int,
|
|
dtype: torch.dtype,
|
|
dp_size: int,
|
|
pp_size: int,
|
|
is_encoder_decoder: bool,
|
|
require_mlp_tp_gather: bool,
|
|
seq_len_fill_value: int,
|
|
encoder_len_fill_value: int,
|
|
num_tokens_per_bs: int,
|
|
cache_loc_dtype: torch.dtype,
|
|
enable_mamba_track: bool,
|
|
ne_token_table: Optional[torch.Tensor] = None,
|
|
hc_hidden_size: Optional[int] = None,
|
|
) -> "DecodeInputBuffers":
|
|
with torch.device(device):
|
|
input_ids = torch.zeros((max_num_token,), dtype=torch.int64)
|
|
input_embeds = torch.zeros((max_num_token, hidden_size), dtype=dtype)
|
|
req_pool_indices = torch.zeros((max_bs,), dtype=torch.int64)
|
|
seq_lens = torch.full((max_bs,), seq_len_fill_value, dtype=torch.int32)
|
|
out_cache_loc = torch.zeros((max_num_token,), dtype=cache_loc_dtype)
|
|
positions = torch.zeros((max_num_token,), dtype=torch.int64)
|
|
mrope_positions = torch.zeros((3, max_num_token), dtype=torch.int64)
|
|
num_token_non_padded = torch.zeros((1,), dtype=torch.int32)
|
|
custom_mask = torch.ones(
|
|
(max_bs * seq_len_fill_value + max_num_token) * num_tokens_per_bs,
|
|
dtype=torch.bool,
|
|
)
|
|
next_token_logits_buffer = torch.zeros(
|
|
(max_num_token, vocab_size),
|
|
dtype=torch.float,
|
|
)
|
|
mamba_track_indices = (
|
|
torch.zeros((max_bs,), dtype=torch.int64)
|
|
if enable_mamba_track
|
|
else None
|
|
)
|
|
mamba_track_mask = (
|
|
torch.zeros((max_bs,), dtype=torch.bool) if enable_mamba_track else None
|
|
)
|
|
|
|
if pp_size > 1:
|
|
# mHC (e.g. DSV4) flattens residual into hidden_states (size = hc_hidden_size).
|
|
is_mhc = hc_hidden_size is not None
|
|
hs = hc_hidden_size if is_mhc else hidden_size
|
|
pp_proxy_tensors = {
|
|
"hidden_states": torch.zeros((max_bs, hs), dtype=dtype),
|
|
}
|
|
if not is_mhc:
|
|
pp_proxy_tensors["residual"] = torch.zeros(
|
|
(max_bs, hidden_size), dtype=dtype
|
|
)
|
|
else:
|
|
pp_proxy_tensors = None
|
|
|
|
if is_encoder_decoder:
|
|
encoder_lens = torch.full(
|
|
(max_bs,), encoder_len_fill_value, dtype=torch.int32
|
|
)
|
|
else:
|
|
encoder_lens = None
|
|
|
|
if require_mlp_tp_gather:
|
|
global_num_tokens_gpu = torch.zeros((dp_size,), dtype=torch.int32)
|
|
global_num_tokens_for_logprob_gpu = torch.zeros(
|
|
(dp_size,), dtype=torch.int32
|
|
)
|
|
else:
|
|
global_num_tokens_gpu = torch.zeros((1,), dtype=torch.int32)
|
|
global_num_tokens_for_logprob_gpu = torch.zeros((1,), dtype=torch.int32)
|
|
|
|
ngram_embedding_info = (
|
|
NgramEmbeddingInfo(
|
|
token_table=ne_token_table,
|
|
column_starts=torch.zeros([max_bs], dtype=torch.int32),
|
|
req_lens=torch.ones([max_bs], dtype=torch.int32),
|
|
out_column_starts=torch.zeros([max_bs], dtype=torch.int32),
|
|
out_req_lens=torch.ones([max_bs], dtype=torch.int32),
|
|
)
|
|
if ne_token_table is not None
|
|
else None
|
|
)
|
|
|
|
# Keep seq_lens_cpu as a true CPU tensor, like the old implementation.
|
|
seq_lens_cpu = torch.full(
|
|
(max_bs,),
|
|
seq_len_fill_value,
|
|
dtype=torch.int32,
|
|
device="cpu",
|
|
)
|
|
|
|
return cls(
|
|
input_ids=input_ids,
|
|
input_embeds=input_embeds,
|
|
req_pool_indices=req_pool_indices,
|
|
seq_lens=seq_lens,
|
|
seq_lens_cpu=seq_lens_cpu,
|
|
out_cache_loc=out_cache_loc,
|
|
positions=positions,
|
|
mrope_positions=mrope_positions,
|
|
num_token_non_padded=num_token_non_padded,
|
|
custom_mask=custom_mask,
|
|
next_token_logits_buffer=next_token_logits_buffer,
|
|
mamba_track_indices=mamba_track_indices,
|
|
mamba_track_mask=mamba_track_mask,
|
|
encoder_lens=encoder_lens,
|
|
global_num_tokens_gpu=global_num_tokens_gpu,
|
|
global_num_tokens_for_logprob_gpu=global_num_tokens_for_logprob_gpu,
|
|
pp_proxy_tensors=pp_proxy_tensors,
|
|
ngram_embedding_info=ngram_embedding_info,
|
|
)
|
|
|
|
def populate_from_forward_batch(
|
|
self,
|
|
*,
|
|
forward_batch: ForwardBatch,
|
|
raw_bs: int,
|
|
raw_num_token: int,
|
|
bs: int,
|
|
seq_len_fill_value: int,
|
|
require_gathered_buffer: bool,
|
|
num_tokens_per_bs: int,
|
|
dsa_enable_prefill_cp: bool,
|
|
enable_num_token_non_padded_flag: bool,
|
|
pp_proxy_tensors: Optional[PPProxyTensors] = None,
|
|
):
|
|
if bs != raw_bs:
|
|
self.seq_lens.fill_(seq_len_fill_value)
|
|
self.out_cache_loc.zero_()
|
|
# Pair with seq_lens fill: padded rows must point at reserved
|
|
# req_pool slot 0 (req_to_token[0, :] is all zeros from init),
|
|
# so dummy attention reads land on slot 0 instead of a stale
|
|
# req_to_token row left by an earlier replay.
|
|
self.req_pool_indices.zero_()
|
|
if self.mamba_track_indices is not None:
|
|
self.mamba_track_indices.zero_()
|
|
if self.mamba_track_mask is not None:
|
|
self.mamba_track_mask.fill_(False)
|
|
|
|
# Build batched copy lists for all GPU tensors.
|
|
dsts = [
|
|
self.input_ids[:raw_num_token],
|
|
self.req_pool_indices[:raw_bs],
|
|
self.seq_lens[:raw_bs],
|
|
self.out_cache_loc[:raw_num_token],
|
|
self.positions[:raw_num_token],
|
|
]
|
|
srcs = [
|
|
forward_batch.input_ids,
|
|
forward_batch.req_pool_indices,
|
|
forward_batch.seq_lens,
|
|
forward_batch.out_cache_loc,
|
|
forward_batch.positions,
|
|
]
|
|
|
|
if self.ngram_embedding_info is not None:
|
|
ngram_embedding_info = forward_batch.ngram_embedding_info
|
|
self.ngram_embedding_info.column_starts[:raw_bs].copy_(
|
|
ngram_embedding_info.column_starts
|
|
)
|
|
self.ngram_embedding_info.req_lens[:raw_bs].copy_(
|
|
ngram_embedding_info.req_lens
|
|
)
|
|
|
|
if (
|
|
self.mamba_track_indices is not None
|
|
and forward_batch.mamba_track_indices is not None
|
|
):
|
|
dsts.append(self.mamba_track_indices[:raw_bs])
|
|
srcs.append(forward_batch.mamba_track_indices)
|
|
if (
|
|
self.mamba_track_mask is not None
|
|
and forward_batch.mamba_track_mask is not None
|
|
):
|
|
dsts.append(self.mamba_track_mask[:raw_bs])
|
|
srcs.append(forward_batch.mamba_track_mask)
|
|
|
|
if self.encoder_lens is not None and forward_batch.encoder_lens is not None:
|
|
dsts.append(self.encoder_lens[:raw_bs])
|
|
srcs.append(forward_batch.encoder_lens)
|
|
|
|
if forward_batch.mrope_positions is not None:
|
|
dsts.append(self.mrope_positions[:, :raw_num_token])
|
|
srcs.append(forward_batch.mrope_positions)
|
|
|
|
if require_gathered_buffer:
|
|
self.global_num_tokens_gpu.fill_(bs * num_tokens_per_bs)
|
|
self.global_num_tokens_for_logprob_gpu.fill_(bs * num_tokens_per_bs)
|
|
|
|
if enable_num_token_non_padded_flag:
|
|
if require_gathered_buffer and not dsa_enable_prefill_cp:
|
|
num_tokens_per_dp = bs * num_tokens_per_bs
|
|
local = compute_local_num_token_non_padded(
|
|
global_num_token_non_padded=forward_batch.num_token_non_padded,
|
|
num_tokens_per_dp=num_tokens_per_dp,
|
|
)
|
|
dsts.append(self.num_token_non_padded)
|
|
srcs.append(local)
|
|
else:
|
|
dsts.append(self.num_token_non_padded)
|
|
srcs.append(forward_batch.num_token_non_padded)
|
|
|
|
# Pipeline-parallel proxy tensors.
|
|
if pp_proxy_tensors is not None and self.pp_proxy_tensors is not None:
|
|
for key, buf in self.pp_proxy_tensors.items():
|
|
src = pp_proxy_tensors.tensors[key]
|
|
dim = src.shape[0]
|
|
dsts.append(buf[:dim])
|
|
srcs.append(src)
|
|
|
|
# Batch all GPU copies, grouped by dtype pair.
|
|
_grouped_foreach_copy_(dsts, srcs)
|
|
|
|
# CPU tensor copy (cannot be batched with GPU tensors).
|
|
if forward_batch.seq_lens_cpu is not None:
|
|
if bs != raw_bs:
|
|
self.seq_lens_cpu.fill_(seq_len_fill_value)
|
|
self.seq_lens_cpu[:raw_bs].copy_(forward_batch.seq_lens_cpu)
|
|
|
|
|
|
# Detect whether the current forward pass is in capture mode
|
|
is_capture_mode = False
|
|
|
|
|
|
def get_is_capture_mode():
|
|
return is_capture_mode
|
|
|
|
|
|
def compile_in_capture_mode(func):
|
|
if get_is_capture_mode():
|
|
return torch.compile(func)
|
|
return func
|
|
|
|
|
|
@contextmanager
|
|
def model_capture_mode():
|
|
global is_capture_mode
|
|
is_capture_mode = True
|
|
|
|
yield
|
|
|
|
is_capture_mode = False
|
|
|
|
|
|
@contextmanager
|
|
def freeze_gc(enable_cudagraph_gc: bool):
|
|
"""
|
|
Optimize garbage collection during CUDA graph capture.
|
|
Clean up, then freeze all remaining objects from being included
|
|
in future collections if GC is disabled during capture.
|
|
"""
|
|
gc.collect()
|
|
should_freeze = not enable_cudagraph_gc
|
|
if should_freeze:
|
|
gc.freeze()
|
|
try:
|
|
yield
|
|
finally:
|
|
if should_freeze:
|
|
gc.unfreeze()
|
|
gc.collect()
|
|
|
|
|
|
def _to_torch(model: torch.nn.Module, reverse: bool, num_tokens: int):
|
|
for sub in model._modules.values():
|
|
if isinstance(sub, MultiPlatformOp):
|
|
if reverse:
|
|
sub.leave_torch_compile()
|
|
else:
|
|
sub.enter_torch_compile(num_tokens=num_tokens)
|
|
if isinstance(sub, torch.nn.Module):
|
|
_to_torch(sub, reverse, num_tokens)
|
|
|
|
|
|
@contextmanager
|
|
def patch_model(
|
|
model: torch.nn.Module,
|
|
enable_compile: bool,
|
|
num_tokens: int,
|
|
tp_group: GroupCoordinator,
|
|
):
|
|
"""Patch the model to make it compatible with with torch.compile"""
|
|
backup_ca_comm = None
|
|
|
|
try:
|
|
if enable_compile:
|
|
_to_torch(model, reverse=False, num_tokens=num_tokens)
|
|
backup_ca_comm = tp_group.ca_comm
|
|
# Use custom-allreduce here.
|
|
# We found the custom allreduce is much faster than the built-in allreduce in torch,
|
|
# even with ENABLE_INTRA_NODE_COMM=1.
|
|
# tp_group.ca_comm = None
|
|
yield torch.compile(
|
|
torch.no_grad()(model.forward),
|
|
mode=os.environ.get(
|
|
"SGLANG_TORCH_COMPILE_MODE", "max-autotune-no-cudagraphs"
|
|
),
|
|
dynamic=_is_hip and get_bool_env_var("SGLANG_TORCH_DYNAMIC_SHAPE"),
|
|
)
|
|
else:
|
|
yield model.forward
|
|
finally:
|
|
if enable_compile:
|
|
_to_torch(model, reverse=True, num_tokens=num_tokens)
|
|
tp_group.ca_comm = backup_ca_comm
|
|
|
|
|
|
def set_torch_compile_config():
|
|
import torch._dynamo.config
|
|
import torch._inductor.config
|
|
|
|
torch._inductor.config.coordinate_descent_tuning = True
|
|
torch._inductor.config.triton.unique_kernel_names = True
|
|
torch._inductor.config.fx_graph_cache = True # Experimental feature to reduce compilation times, will be on by default in future
|
|
|
|
# FIXME: tmp workaround
|
|
torch._dynamo.config.accumulated_cache_size_limit = 1024
|
|
if hasattr(torch._dynamo.config, "cache_size_limit"):
|
|
torch._dynamo.config.cache_size_limit = 1024
|
|
|
|
monkey_patch_torch_compile()
|
|
|
|
|
|
def get_batch_sizes_to_capture(model_runner: ModelRunner, num_tokens_per_bs=1):
|
|
server_args = model_runner.server_args
|
|
capture_bs = server_args.cuda_graph_bs
|
|
num_max_requests = model_runner.req_to_token_pool.size
|
|
|
|
mul_base = 1
|
|
if server_args.enable_two_batch_overlap:
|
|
mul_base *= 2
|
|
num_tokens_per_bs = 1 # tbo not test, set num_tokens_per_bs to 1
|
|
|
|
if require_gathered_buffer(server_args):
|
|
mul_base *= get_attention_tp_size()
|
|
|
|
if mul_base % get_attention_cp_size() != 0:
|
|
mul_base *= get_attention_cp_size()
|
|
|
|
# pad `num_max_requests` to avoid being filtered out
|
|
num_max_requests = (num_max_requests + mul_base - 1) // mul_base * mul_base
|
|
if max(capture_bs) > num_max_requests:
|
|
# In some cases (e.g., with a small GPU or --max-running-requests), the #max-running-requests
|
|
# is very small. We add more values here to make sure we capture the maximum bs.
|
|
capture_bs += [num_max_requests]
|
|
|
|
# Model input token count = bs * num_tokens_per_bs; must be a multiple of attn_tp_size.
|
|
capture_bs = [bs for bs in capture_bs if bs * num_tokens_per_bs % mul_base == 0]
|
|
capture_bs = [bs for bs in capture_bs if bs <= num_max_requests]
|
|
capture_bs = list(sorted(set(capture_bs)))
|
|
|
|
assert len(capture_bs) > 0 and capture_bs[0] > 0, f"{capture_bs=}"
|
|
compile_bs = (
|
|
[bs for bs in capture_bs if bs <= server_args.torch_compile_max_bs]
|
|
if server_args.enable_torch_compile
|
|
else []
|
|
)
|
|
return capture_bs, compile_bs
|
|
|
|
|
|
# Reuse this memory pool across all cuda graph runners.
|
|
global_graph_memory_pool = None
|
|
|
|
|
|
def get_global_graph_memory_pool():
|
|
return global_graph_memory_pool
|
|
|
|
|
|
def set_global_graph_memory_pool(val):
|
|
global global_graph_memory_pool
|
|
global_graph_memory_pool = val
|
|
|
|
|
|
class CudaGraphRunner:
|
|
"""A CudaGraphRunner runs the forward pass of a model with cuda graph and torch.compile."""
|
|
|
|
def __init__(
|
|
self,
|
|
model_runner: ModelRunner,
|
|
*,
|
|
attn_backend=None,
|
|
speculative_num_steps: Optional[int] = None,
|
|
speculative_num_draft_tokens: Optional[int] = None,
|
|
):
|
|
# Parse args
|
|
self.model_runner = model_runner
|
|
self.device = model_runner.device
|
|
self.device_module = torch.get_device_module(self.device)
|
|
self.graphs = {}
|
|
self.output_buffers = {}
|
|
self.enable_torch_compile = model_runner.server_args.enable_torch_compile
|
|
self.disable_padding = model_runner.server_args.disable_cuda_graph_padding
|
|
self.is_encoder_decoder = model_runner.model_config.is_encoder_decoder
|
|
self.require_gathered_buffer = require_gathered_buffer(model_runner.server_args)
|
|
self.require_mlp_tp_gather = require_mlp_tp_gather(model_runner.server_args)
|
|
self.require_mlp_sync = require_mlp_sync(model_runner.server_args)
|
|
self.require_attn_tp_gather = require_attn_tp_gather(model_runner.server_args)
|
|
self.enable_two_batch_overlap = (
|
|
model_runner.server_args.enable_two_batch_overlap
|
|
)
|
|
self.use_ngram_embedding = model_runner.use_ngram_embedding
|
|
if self.use_ngram_embedding:
|
|
hf_config = model_runner.model_config.hf_config
|
|
self.ngram_embedding_n = hf_config.ngram_embedding_n
|
|
self.ngram_embedding_k = hf_config.ngram_embedding_k
|
|
self.speculative_algorithm = model_runner.server_args.speculative_algorithm
|
|
self.enable_profile_cuda_graph = (
|
|
model_runner.server_args.enable_profile_cuda_graph
|
|
)
|
|
self.tp_size = model_runner.server_args.tp_size
|
|
self.dp_size = model_runner.server_args.dp_size
|
|
self.pp_size = model_runner.server_args.pp_size
|
|
self.enable_pdmux = model_runner.server_args.enable_pdmux
|
|
|
|
self.attn_tp_size = get_attention_tp_size()
|
|
self.attn_tp_rank = get_attention_tp_rank()
|
|
# True if a DSACPLayerCommunicator-style prefill-CP flavor is active
|
|
# (DSA or MLA). These flavors feed a zigzag-split rank-local layout
|
|
# into the runner; MHA-arch prefill CP (Qwen3/Qwen2 MoE via PR
|
|
# #18233) uses the plain LayerCommunicator with an attn_tp-replicated
|
|
# layout and is intentionally excluded so the attn_tp-local
|
|
# num_token_non_padded adjustment still runs for it.
|
|
self.enable_prefill_cp = (
|
|
is_dsa_enable_prefill_cp() or is_mla_prefill_cp_enabled()
|
|
)
|
|
|
|
self.deepep_adapter = DeepEPCudaGraphRunnerAdapter()
|
|
|
|
self.dllm_config = DllmConfig.from_server_args(model_runner.server_args)
|
|
self.is_dllm = self.dllm_config is not None
|
|
self.attn_backend = attn_backend or model_runner.attn_backend
|
|
self.speculative_num_steps = (
|
|
model_runner.server_args.speculative_num_steps
|
|
if speculative_num_steps is None
|
|
else speculative_num_steps
|
|
)
|
|
self.speculative_num_draft_tokens = (
|
|
model_runner.server_args.speculative_num_draft_tokens
|
|
if speculative_num_draft_tokens is None
|
|
else speculative_num_draft_tokens
|
|
)
|
|
|
|
self.capture_forward_mode = ForwardMode.DECODE
|
|
self.capture_hidden_mode = CaptureHiddenMode.NULL
|
|
self.num_tokens_per_bs = 1
|
|
if model_runner.spec_algorithm.is_speculative():
|
|
if self.model_runner.is_draft_worker:
|
|
# Draft workers can use TARGET_VERIFY mode.
|
|
if (
|
|
not self.model_runner.spec_algorithm.supports_target_verify_for_draft()
|
|
):
|
|
raise RuntimeError("This should not happen")
|
|
self.capture_forward_mode = ForwardMode.TARGET_VERIFY
|
|
self.num_tokens_per_bs = (
|
|
model_runner.spec_algorithm.get_num_tokens_per_bs_for_target_verify(
|
|
self.speculative_num_draft_tokens, model_runner.is_draft_worker
|
|
)
|
|
)
|
|
elif self.is_dllm:
|
|
self.capture_forward_mode = ForwardMode.DLLM_EXTEND
|
|
self.num_tokens_per_bs = self.dllm_config.block_size
|
|
|
|
# Batch sizes to capture
|
|
self.capture_bs, self.compile_bs = get_batch_sizes_to_capture(
|
|
model_runner, self.num_tokens_per_bs
|
|
)
|
|
log_info_on_rank0(logger, f"Capture cuda graph bs {self.capture_bs}")
|
|
if KTRANSFORMERS_AVAILABLE:
|
|
KTMoEWrapper.set_capture_batch_sizes(self.capture_bs)
|
|
|
|
# If returning hidden states is enabled, set initial capture hidden mode to full to avoid double-capture on startup
|
|
if model_runner.server_args.enable_return_hidden_states:
|
|
self.capture_hidden_mode = CaptureHiddenMode.FULL
|
|
|
|
# Attention backend
|
|
self.max_bs = max(self.capture_bs)
|
|
self.max_num_token = self.max_bs * self.num_tokens_per_bs
|
|
self.attn_backend.init_cuda_graph_state(self.max_bs, self.max_num_token)
|
|
|
|
# Init PDMux if needed
|
|
self.maybe_init_pdmux()
|
|
self.seq_len_fill_value = (
|
|
self.attn_backend.get_cuda_graph_seq_len_fill_value()
|
|
if self.dllm_config is None
|
|
else self.dllm_config.block_size
|
|
)
|
|
|
|
# Non-zero encoder length ensures cross-attention kernels are captured in the graph.
|
|
self.encoder_len_fill_value = (
|
|
getattr(model_runner.model_config.hf_config, "max_source_positions", 0)
|
|
if self.is_encoder_decoder
|
|
else 0
|
|
)
|
|
|
|
if self.enable_torch_compile:
|
|
set_torch_compile_config()
|
|
|
|
if self.model_runner.server_args.enable_lora:
|
|
# Phase 2 of LoRA CUDA graph init: dense LoRA batch metadata.
|
|
# Phase 1 (MoE buffers) was handled earlier in ModelRunner via
|
|
# lora_manager.init_cuda_graph_moe_buffers().
|
|
self.model_runner.lora_manager.init_cuda_graph_batch_info(
|
|
max_bs_in_cuda_graph=self.max_bs,
|
|
num_tokens_per_bs=self.num_tokens_per_bs,
|
|
)
|
|
|
|
enable_mamba_track = (
|
|
self.model_runner.server_args.enable_mamba_extra_buffer()
|
|
and self.model_runner.spec_algorithm.is_none()
|
|
)
|
|
|
|
if self.require_gathered_buffer:
|
|
assert self.require_mlp_tp_gather or self.require_attn_tp_gather
|
|
self.buffers: DecodeInputBuffers = DecodeInputBuffers.create(
|
|
device=self.device,
|
|
max_bs=self.max_bs,
|
|
max_num_token=self.max_num_token,
|
|
hidden_size=self.model_runner.model_config.hidden_size,
|
|
vocab_size=self.model_runner.model_config.vocab_size,
|
|
dtype=self.model_runner.model_config.dtype,
|
|
dp_size=self.dp_size,
|
|
pp_size=self.pp_size,
|
|
is_encoder_decoder=self.is_encoder_decoder,
|
|
require_mlp_tp_gather=self.require_mlp_tp_gather,
|
|
seq_len_fill_value=self.seq_len_fill_value,
|
|
encoder_len_fill_value=self.encoder_len_fill_value,
|
|
num_tokens_per_bs=self.num_tokens_per_bs,
|
|
cache_loc_dtype=self._cache_loc_dtype(),
|
|
enable_mamba_track=enable_mamba_track,
|
|
ne_token_table=(
|
|
model_runner.token_table if self.use_ngram_embedding else None
|
|
),
|
|
hc_hidden_size=getattr(
|
|
self.model_runner.model_config, "hc_hidden_size", None
|
|
),
|
|
)
|
|
self.buffers.share_buffers()
|
|
|
|
self.tbo_plugin = TboCudaGraphRunnerPlugin()
|
|
|
|
# Capture
|
|
try:
|
|
with model_capture_mode():
|
|
self.capture()
|
|
except RuntimeError as e:
|
|
raise Exception(
|
|
f"Capture cuda graph failed: {e}\n{CUDA_GRAPH_CAPTURE_FAILED_MSG}"
|
|
)
|
|
|
|
def maybe_init_pdmux(self):
|
|
if self.enable_pdmux:
|
|
self.stream_groups = get_stream_groups()
|
|
for attn_backend in self.model_runner.decode_attn_backend_group:
|
|
attn_backend.init_cuda_graph_state(self.max_bs, self.max_num_token)
|
|
|
|
def _cache_loc_dtype(self):
|
|
return torch.int64
|
|
|
|
def can_run(self, forward_batch: ForwardBatch):
|
|
# Disable for token embedding overrides (dynamic per-request)
|
|
if forward_batch.replace_embeds is not None:
|
|
return False
|
|
if self.require_mlp_tp_gather:
|
|
cuda_graph_bs = (
|
|
max(forward_batch.global_num_tokens_cpu) // self.num_tokens_per_bs
|
|
if self.model_runner.spec_algorithm.is_eagle()
|
|
or self.model_runner.spec_algorithm.is_standalone()
|
|
or self.model_runner.spec_algorithm.is_dflash()
|
|
else max(forward_batch.global_num_tokens_cpu)
|
|
)
|
|
else:
|
|
cuda_graph_bs = forward_batch.batch_size
|
|
|
|
graph_key = cuda_graph_bs
|
|
if self.enable_pdmux:
|
|
graph_key = f"{get_current_stream_idx()}_{cuda_graph_bs}"
|
|
|
|
is_bs_supported = (
|
|
graph_key in self.graphs
|
|
if self.disable_padding
|
|
else cuda_graph_bs <= self.max_bs
|
|
)
|
|
|
|
if self.require_mlp_sync:
|
|
is_bs_supported = is_bs_supported and forward_batch.can_run_dp_cuda_graph
|
|
|
|
# NOTE: cuda graph cannot handle mixed batch (encoder_len = 0)
|
|
# If mixed batch cannot be supported, then encoder_lens can be removed in cuda graph
|
|
# because the full_text_row_masked_out_mask tensor will always be ones
|
|
is_encoder_lens_supported = (
|
|
torch.all(forward_batch.encoder_lens > 0)
|
|
if self.is_encoder_decoder
|
|
else True
|
|
)
|
|
|
|
requested_capture_hidden_mode = max(
|
|
forward_batch.capture_hidden_mode,
|
|
(
|
|
forward_batch.spec_info.capture_hidden_mode
|
|
if getattr(forward_batch.spec_info, "capture_hidden_mode", None)
|
|
is not None
|
|
else CaptureHiddenMode.NULL
|
|
),
|
|
)
|
|
capture_hidden_mode_matches = (
|
|
requested_capture_hidden_mode == CaptureHiddenMode.NULL
|
|
or requested_capture_hidden_mode == self.capture_hidden_mode
|
|
)
|
|
is_tbo_supported = (
|
|
forward_batch.can_run_tbo if self.enable_two_batch_overlap else True
|
|
)
|
|
|
|
is_ngram_supported = (
|
|
(
|
|
forward_batch.batch_size * self.num_tokens_per_bs
|
|
== forward_batch.input_ids.numel()
|
|
)
|
|
if self.model_runner.spec_algorithm.is_ngram()
|
|
else True
|
|
)
|
|
|
|
return (
|
|
is_bs_supported
|
|
and is_encoder_lens_supported
|
|
and is_tbo_supported
|
|
and capture_hidden_mode_matches
|
|
and is_ngram_supported
|
|
)
|
|
|
|
def _init_profile_context_and_memory_record(self):
|
|
profile_context = profile(
|
|
activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
|
|
record_shapes=True,
|
|
)
|
|
torch.cuda.memory._record_memory_history()
|
|
return profile_context
|
|
|
|
def _post_process_after_profile(self, prof_context):
|
|
torch.cuda.memory._dump_snapshot(f"cuda_graph_runner_memory_usage.pickle")
|
|
torch.cuda.memory._record_memory_history(enabled=None)
|
|
log_message = (
|
|
"Sorted by CUDA Time:\n"
|
|
+ prof_context.key_averages(group_by_input_shape=True).table(
|
|
sort_by="cuda_time_total", row_limit=10
|
|
)
|
|
+ "\n\nSorted by CPU Time:\n"
|
|
+ prof_context.key_averages(group_by_input_shape=True).table(
|
|
sort_by="cpu_time_total", row_limit=10
|
|
)
|
|
+ "\n\nMemory Usage is saved to cuda_graph_runner_memory_usage.pickle\n"
|
|
)
|
|
logger.info(log_message)
|
|
|
|
def capture(self) -> None:
|
|
profile_context = empty_context()
|
|
if self.enable_profile_cuda_graph:
|
|
profile_context = self._init_profile_context_and_memory_record()
|
|
|
|
def _capture_one_stream(stream_idx: Optional[int] = None):
|
|
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_bs)))
|
|
if get_tensor_model_parallel_rank() == 0
|
|
else reversed(self.capture_bs)
|
|
)
|
|
for i, bs 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 batches ({bs=} {avail_mem=:.2f} GB)"
|
|
)
|
|
|
|
with patch_model(
|
|
self.model_runner.model,
|
|
bs in self.compile_bs,
|
|
num_tokens=bs * self.num_tokens_per_bs,
|
|
tp_group=self.model_runner.tp_group,
|
|
) as forward:
|
|
(
|
|
graph,
|
|
output_buffers,
|
|
) = self.capture_one_batch_size(bs, forward, stream_idx)
|
|
# For pd_multiplexing, we need to save the graph and output buffers
|
|
key = bs if stream_idx is None else f"{stream_idx}_{bs}"
|
|
self.graphs[key] = graph
|
|
self.output_buffers[key] = output_buffers
|
|
|
|
# Trigger CUDA graph capture for specific shapes.
|
|
# Capture the large shapes first so that the smaller shapes
|
|
# can reuse the memory pool allocated for the large shapes.
|
|
with freeze_gc(self.model_runner.server_args.enable_cudagraph_gc):
|
|
if not self.enable_pdmux:
|
|
with graph_capture() as graph_capture_context, profile_context as prof:
|
|
self.stream = graph_capture_context.stream
|
|
_capture_one_stream()
|
|
else:
|
|
set_pdmux_status(False)
|
|
for i, sg in enumerate(self.stream_groups):
|
|
with (
|
|
graph_capture(stream=sg[1]) as graph_capture_context,
|
|
profile_context as prof,
|
|
):
|
|
self.stream = graph_capture_context.stream
|
|
_capture_one_stream(i)
|
|
|
|
if self.enable_profile_cuda_graph:
|
|
self._post_process_after_profile(prof)
|
|
|
|
def _capture_graph(self, graph, pool, stream, run_once_fn):
|
|
if self.model_runner.server_args.debug_cuda_graph:
|
|
assert (
|
|
envs.SGLANG_USE_BREAKABLE_CUDA_GRAPH.get()
|
|
), "Breakable CUDA graph is not enabled in debug mode"
|
|
|
|
memory_saver_adapter = TorchMemorySaverAdapter.create(
|
|
enable=self.model_runner.server_args.enable_memory_saver
|
|
and get_bool_env_var("SGLANG_MEMORY_SAVER_CUDA_GRAPH")
|
|
)
|
|
|
|
if envs.SGLANG_USE_BREAKABLE_CUDA_GRAPH.get():
|
|
if memory_saver_adapter.enabled:
|
|
raise NotImplementedError(
|
|
"Breakable CUDA graph is not compatible with memory saver mode"
|
|
)
|
|
graph_ctx = BreakableCUDAGraphCapture
|
|
else:
|
|
graph_ctx = (
|
|
partial(memory_saver_adapter.cuda_graph, tag=GPU_MEMORY_TYPE_CUDA_GRAPH)
|
|
if memory_saver_adapter.enabled
|
|
else self.device_module.graph
|
|
)
|
|
|
|
if self.model_runner.server_args.debug_cuda_graph:
|
|
captured_fn = eager_on_graph(True)(run_once_fn)
|
|
else:
|
|
captured_fn = run_once_fn
|
|
|
|
with graph_ctx(cuda_graph=graph, pool=pool, stream=stream):
|
|
out = captured_fn()
|
|
return out
|
|
|
|
def _create_device_graph(self):
|
|
if envs.SGLANG_USE_BREAKABLE_CUDA_GRAPH.get():
|
|
if _is_hip:
|
|
raise RuntimeError("Breakable CUDA graph is not supported on ROCm/HIP")
|
|
return BreakableCUDAGraph()
|
|
return torch.cuda.CUDAGraph()
|
|
|
|
def capture_one_batch_size(
|
|
self, bs: int, forward: Callable, stream_idx: Optional[int] = None
|
|
):
|
|
buffers: DecodeInputBuffers = self.buffers
|
|
graph = self._create_device_graph()
|
|
stream = self.stream
|
|
num_tokens = bs * self.num_tokens_per_bs
|
|
|
|
# Graph inputs
|
|
input_ids = buffers.input_ids[:num_tokens]
|
|
req_pool_indices = buffers.req_pool_indices[:bs]
|
|
seq_lens = buffers.seq_lens[:bs]
|
|
seq_lens_cpu = buffers.seq_lens_cpu[:bs]
|
|
out_cache_loc = buffers.out_cache_loc[:num_tokens]
|
|
positions = buffers.positions[:num_tokens]
|
|
if self.is_encoder_decoder:
|
|
encoder_lens = buffers.encoder_lens[:bs]
|
|
else:
|
|
encoder_lens = None
|
|
mrope_positions = buffers.mrope_positions[:, :num_tokens]
|
|
next_token_logits_buffer = buffers.next_token_logits_buffer[:num_tokens]
|
|
|
|
# Adjust for attention TP if needed (matching replay path in
|
|
# populate_from_forward_batch).
|
|
buffers.num_token_non_padded[...] = num_tokens
|
|
if (
|
|
enable_num_token_non_padded()
|
|
and self.require_gathered_buffer
|
|
and not self.enable_prefill_cp
|
|
):
|
|
local = compute_local_num_token_non_padded(
|
|
global_num_token_non_padded=buffers.num_token_non_padded,
|
|
num_tokens_per_dp=num_tokens,
|
|
)
|
|
buffers.num_token_non_padded.copy_(local)
|
|
|
|
# pipeline parallelism
|
|
if self.pp_size > 1:
|
|
pp_proxy_tensors = PPProxyTensors(
|
|
{k: v[:num_tokens] for k, v in buffers.pp_proxy_tensors.items()}
|
|
)
|
|
|
|
if self.require_mlp_tp_gather:
|
|
buffers.global_num_tokens_gpu.copy_(
|
|
torch.tensor(
|
|
[num_tokens] * self.dp_size,
|
|
dtype=torch.int32,
|
|
device=input_ids.device,
|
|
)
|
|
)
|
|
buffers.global_num_tokens_for_logprob_gpu.copy_(
|
|
torch.tensor(
|
|
[num_tokens] * self.dp_size,
|
|
dtype=torch.int32,
|
|
device=input_ids.device,
|
|
)
|
|
)
|
|
global_dp_buffer_len = num_tokens * self.dp_size
|
|
elif self.require_attn_tp_gather:
|
|
buffers.global_num_tokens_gpu.copy_(
|
|
torch.tensor(
|
|
[num_tokens],
|
|
dtype=torch.int32,
|
|
device=input_ids.device,
|
|
)
|
|
)
|
|
buffers.global_num_tokens_for_logprob_gpu.copy_(
|
|
torch.tensor(
|
|
[num_tokens],
|
|
dtype=torch.int32,
|
|
device=input_ids.device,
|
|
)
|
|
)
|
|
global_dp_buffer_len = num_tokens
|
|
else:
|
|
global_dp_buffer_len = None
|
|
|
|
spec_info = self.get_spec_info(num_tokens)
|
|
if self.capture_hidden_mode != CaptureHiddenMode.FULL:
|
|
self.capture_hidden_mode = (
|
|
spec_info.capture_hidden_mode if spec_info else CaptureHiddenMode.NULL
|
|
)
|
|
|
|
if self.model_runner.server_args.enable_lora:
|
|
# It is safe to capture CUDA graph using empty LoRA id, as the LoRA kernels will always be launched whenever
|
|
# `--enable-lora` is set to True (and return immediately if the LoRA id is empty for perf optimization).
|
|
lora_ids = [None] * bs
|
|
else:
|
|
lora_ids = None
|
|
|
|
# mamba state tracking
|
|
mamba_track_indices = (
|
|
buffers.mamba_track_indices[:bs]
|
|
if buffers.mamba_track_indices is not None
|
|
else None
|
|
)
|
|
mamba_track_mask = (
|
|
buffers.mamba_track_mask[:bs]
|
|
if buffers.mamba_track_mask is not None
|
|
else None
|
|
)
|
|
|
|
if stream_idx is None:
|
|
attn_backend = self.attn_backend
|
|
else:
|
|
assert self.enable_pdmux
|
|
attn_backend = self.model_runner.decode_attn_backend_group[stream_idx]
|
|
|
|
forward_batch = ForwardBatch(
|
|
forward_mode=self.capture_forward_mode,
|
|
batch_size=bs,
|
|
input_ids=input_ids,
|
|
req_pool_indices=req_pool_indices,
|
|
seq_lens=seq_lens,
|
|
seq_lens_cpu=seq_lens_cpu,
|
|
next_token_logits_buffer=next_token_logits_buffer,
|
|
orig_seq_lens=seq_lens,
|
|
out_cache_loc=out_cache_loc,
|
|
seq_lens_sum=seq_lens.sum().item(),
|
|
mamba_track_indices=mamba_track_indices,
|
|
mamba_track_mask=mamba_track_mask,
|
|
mamba_track_seqlens=None, # Prefill only
|
|
encoder_lens=encoder_lens,
|
|
return_logprob=False,
|
|
positions=positions,
|
|
global_num_tokens_gpu=buffers.global_num_tokens_gpu,
|
|
global_num_tokens_for_logprob_gpu=buffers.global_num_tokens_for_logprob_gpu,
|
|
dp_padding_mode=DpPaddingMode.get_default_mode_in_cuda_graph(),
|
|
global_dp_buffer_len=global_dp_buffer_len,
|
|
mrope_positions=mrope_positions,
|
|
spec_algorithm=self.model_runner.spec_algorithm,
|
|
spec_info=spec_info,
|
|
capture_hidden_mode=self.capture_hidden_mode,
|
|
num_token_non_padded=buffers.num_token_non_padded,
|
|
global_forward_mode=self.capture_forward_mode,
|
|
lora_ids=lora_ids,
|
|
)
|
|
|
|
# Trip the coordinator so the hisparse code path is captured into the
|
|
# graph; backends read it from self.model_runner.hisparse_coordinator.
|
|
hisparse_coordinator = self.model_runner.hisparse_coordinator
|
|
if hisparse_coordinator is not None:
|
|
hisparse_coordinator.num_real_reqs.fill_(bs)
|
|
|
|
if buffers.ngram_embedding_info is not None:
|
|
forward_batch.ngram_embedding_info = buffers.ngram_embedding_info.slice(bs)
|
|
|
|
# All setup hooks below read get_attn_backend() (TboForwardBatchPreparer,
|
|
# DeepEP adapter, …) so they must run inside the same ForwardContext
|
|
# that wraps the warmup/capture forward.
|
|
with forward_context(ForwardContext(attn_backend=attn_backend)):
|
|
self.tbo_plugin.capture_one_batch_size(forward_batch, num_tokens=num_tokens)
|
|
|
|
if lora_ids is not None:
|
|
self.model_runner.lora_manager.prepare_lora_batch(forward_batch)
|
|
|
|
attn_backend.init_forward_metadata_capture_cuda_graph(
|
|
bs,
|
|
num_tokens,
|
|
req_pool_indices,
|
|
seq_lens,
|
|
encoder_lens,
|
|
forward_batch.forward_mode,
|
|
forward_batch.spec_info,
|
|
)
|
|
|
|
def run_once():
|
|
# Without this, warmup-1 caches the translation; the capture
|
|
# run hits the cache, skips the gather, and replay reuses
|
|
# stale SWA locations.
|
|
if self.model_runner.is_hybrid_swa:
|
|
self.model_runner.token_to_kv_pool.invalidate_loc_cache()
|
|
|
|
forward_batch.dp_local_start_pos = forward_batch.dp_local_num_tokens = (
|
|
None
|
|
)
|
|
set_dp_buffer_len(
|
|
global_dp_buffer_len,
|
|
num_tokens,
|
|
forward_batch.dp_padding_mode.is_max_len(),
|
|
)
|
|
set_is_extend_in_batch(False)
|
|
|
|
kwargs = {}
|
|
if (
|
|
self.pp_size > 1
|
|
and "pp_proxy_tensors" in inspect.signature(forward).parameters
|
|
):
|
|
kwargs["pp_proxy_tensors"] = PPProxyTensors(
|
|
{k: v.clone() for k, v in pp_proxy_tensors.tensors.items()}
|
|
)
|
|
if (
|
|
self.model_runner.spec_algorithm.is_dflash()
|
|
and self.model_runner.is_draft_worker
|
|
and "input_embeds" in inspect.signature(forward).parameters
|
|
):
|
|
kwargs["input_embeds"] = buffers.input_embeds[:num_tokens]
|
|
|
|
logits_output_or_pp_proxy_tensors = forward(
|
|
input_ids,
|
|
forward_batch.positions,
|
|
forward_batch,
|
|
**kwargs,
|
|
)
|
|
return logits_output_or_pp_proxy_tensors
|
|
|
|
self.deepep_adapter.capture(is_extend_in_batch=False)
|
|
|
|
canary_ctx = contextlib.nullcontext()
|
|
with canary_ctx:
|
|
for _ in range(2):
|
|
self.device_module.synchronize()
|
|
self.model_runner.tp_group.barrier()
|
|
run_once()
|
|
attn_backend.on_after_cuda_graph_warmup()
|
|
|
|
if get_global_graph_memory_pool() is None:
|
|
set_global_graph_memory_pool(self.device_module.graph_pool_handle())
|
|
# Set graph pool id globally to be able to use symmetric memory
|
|
set_graph_pool_id(get_global_graph_memory_pool())
|
|
|
|
out = self._capture_graph(
|
|
graph, get_global_graph_memory_pool(), stream, run_once
|
|
)
|
|
|
|
return graph, out
|
|
|
|
def recapture_if_needed(self, forward_batch: ForwardBatch):
|
|
|
|
# If the required capture_hidden_mode changes, we need to recapture the graph
|
|
|
|
# These are the different factors that can influence the capture_hidden_mode
|
|
capture_hidden_mode_required_by_forward_batch = (
|
|
forward_batch.capture_hidden_mode
|
|
)
|
|
capture_hidden_mode_required_by_spec_info = (
|
|
getattr(forward_batch.spec_info, "capture_hidden_mode", None)
|
|
or CaptureHiddenMode.NULL
|
|
)
|
|
capture_hidden_mode_required_for_returning_hidden_states = (
|
|
CaptureHiddenMode.FULL
|
|
if self.model_runner.server_args.enable_return_hidden_states
|
|
else CaptureHiddenMode.NULL
|
|
)
|
|
|
|
# Determine the highest capture_hidden_mode required
|
|
# (If we have FULL, we can emulate LAST or NULL)
|
|
# (If we have LAST, we can emulate NULL)
|
|
required_capture_hidden_mode = max(
|
|
capture_hidden_mode_required_by_forward_batch,
|
|
capture_hidden_mode_required_by_spec_info,
|
|
capture_hidden_mode_required_for_returning_hidden_states,
|
|
)
|
|
|
|
# If the current hidden mode is no longer aligned with the required hidden mode, we need to set it to what is required and re-capture
|
|
if self.capture_hidden_mode != required_capture_hidden_mode:
|
|
self.capture_hidden_mode = required_capture_hidden_mode
|
|
self.capture()
|
|
|
|
def replay_prepare(
|
|
self,
|
|
forward_batch: ForwardBatch,
|
|
pp_proxy_tensors: Optional[PPProxyTensors] = None,
|
|
):
|
|
buffers = self.buffers
|
|
self.recapture_if_needed(forward_batch)
|
|
|
|
raw_bs = forward_batch.batch_size
|
|
raw_num_token = raw_bs * self.num_tokens_per_bs
|
|
|
|
# Pad
|
|
if self.require_mlp_tp_gather:
|
|
max_num_tokens = max(forward_batch.global_num_tokens_cpu)
|
|
max_batch_size = (
|
|
max_num_tokens / self.num_tokens_per_bs
|
|
if self.model_runner.spec_algorithm.is_eagle()
|
|
or self.model_runner.spec_algorithm.is_standalone()
|
|
or self.model_runner.spec_algorithm.is_dflash()
|
|
else max_num_tokens
|
|
)
|
|
index = bisect.bisect_left(self.capture_bs, max_batch_size)
|
|
else:
|
|
index = bisect.bisect_left(self.capture_bs, raw_bs)
|
|
bs = self.capture_bs[index]
|
|
|
|
buffers.populate_from_forward_batch(
|
|
forward_batch=forward_batch,
|
|
raw_bs=raw_bs,
|
|
raw_num_token=raw_num_token,
|
|
bs=bs,
|
|
seq_len_fill_value=self.seq_len_fill_value,
|
|
require_gathered_buffer=self.require_gathered_buffer,
|
|
num_tokens_per_bs=self.num_tokens_per_bs,
|
|
# Parameter name retained for API stability; semantically this is
|
|
# "any prefill-CP flavor enabled" (DSA CP or MLA CP).
|
|
dsa_enable_prefill_cp=self.enable_prefill_cp,
|
|
enable_num_token_non_padded_flag=enable_num_token_non_padded(),
|
|
pp_proxy_tensors=pp_proxy_tensors,
|
|
)
|
|
|
|
if (
|
|
self.model_runner.spec_algorithm.is_dflash()
|
|
and self.model_runner.is_draft_worker
|
|
and forward_batch.input_embeds is not None
|
|
):
|
|
buffers.input_embeds[:raw_num_token].copy_(forward_batch.input_embeds)
|
|
# Padded tokens aren't read, so skip zeroing them.
|
|
if self.enable_two_batch_overlap:
|
|
self.tbo_plugin.replay_prepare(
|
|
forward_mode=self.capture_forward_mode,
|
|
bs=bs,
|
|
num_token_non_padded=len(forward_batch.input_ids),
|
|
spec_info=forward_batch.spec_info,
|
|
)
|
|
if forward_batch.forward_mode.is_idle() and forward_batch.spec_info is not None:
|
|
forward_batch.spec_info.custom_mask = buffers.custom_mask
|
|
# Attention backend
|
|
if self.enable_pdmux:
|
|
stream_idx = get_current_stream_idx()
|
|
attn_backend = self.model_runner.decode_attn_backend_group[stream_idx]
|
|
else:
|
|
attn_backend = self.attn_backend
|
|
# FIXME: implicit channel for backends (dsv4) that need forward_batch
|
|
# in replay metadata prep. Should become a real param on the interface.
|
|
attn_backend._replay_forward_batch = forward_batch
|
|
seq_lens_sum_arg = (
|
|
None
|
|
if forward_batch.seq_lens_sum is None
|
|
else forward_batch.seq_lens_sum + (bs - raw_bs) * self.seq_len_fill_value
|
|
)
|
|
attn_backend.init_forward_metadata_replay_cuda_graph(
|
|
bs,
|
|
buffers.req_pool_indices[:bs],
|
|
buffers.seq_lens[:bs],
|
|
seq_lens_sum_arg,
|
|
buffers.encoder_lens[:bs] if self.is_encoder_decoder else None,
|
|
self.capture_forward_mode,
|
|
forward_batch.spec_info,
|
|
seq_lens_cpu=buffers.seq_lens_cpu[:bs],
|
|
)
|
|
attn_backend._replay_forward_batch = None
|
|
|
|
# Store fields
|
|
self.raw_bs = raw_bs
|
|
self.raw_num_token = raw_num_token
|
|
self.bs = bs
|
|
|
|
if self.model_runner.hisparse_coordinator is not None:
|
|
self.model_runner.hisparse_coordinator.num_real_reqs.fill_(raw_bs)
|
|
|
|
def replay(
|
|
self,
|
|
forward_batch: ForwardBatch,
|
|
skip_attn_backend_init: bool = False,
|
|
pp_proxy_tensors: Optional[PPProxyTensors] = None,
|
|
) -> Union[LogitsProcessorOutput, PPProxyTensors]:
|
|
self.deepep_adapter.replay()
|
|
|
|
if not skip_attn_backend_init:
|
|
self.replay_prepare(forward_batch, pp_proxy_tensors)
|
|
else:
|
|
# In speculative decoding, these two fields are still needed.
|
|
self.buffers.input_ids[: self.raw_num_token].copy_(forward_batch.input_ids)
|
|
self.buffers.positions[: self.raw_num_token].copy_(forward_batch.positions)
|
|
if (
|
|
self.model_runner.spec_algorithm.is_dflash()
|
|
and self.model_runner.is_draft_worker
|
|
and forward_batch.input_embeds is not None
|
|
):
|
|
self.buffers.input_embeds[: self.raw_num_token].copy_(
|
|
forward_batch.input_embeds
|
|
)
|
|
|
|
# Replay
|
|
if self.enable_pdmux:
|
|
graph_key = f"{get_current_stream_idx()}_{self.bs}"
|
|
else:
|
|
graph_key = self.bs
|
|
ctx = (
|
|
self.model_runner.device_timer.wrap(
|
|
metadata={
|
|
"category": forward_batch.forward_mode.name.lower(),
|
|
}
|
|
)
|
|
if self.model_runner.device_timer
|
|
else contextlib.nullcontext()
|
|
)
|
|
with ctx:
|
|
self.graphs[graph_key].replay()
|
|
|
|
output = self.output_buffers[graph_key]
|
|
|
|
if isinstance(output, LogitsProcessorOutput):
|
|
if self.is_dllm:
|
|
next_token_logits = None
|
|
full_logits = (
|
|
output.full_logits[: self.raw_num_token]
|
|
if output.full_logits is not None
|
|
else None
|
|
)
|
|
else:
|
|
full_logits = None
|
|
next_token_logits = (
|
|
output.next_token_logits[: self.raw_num_token]
|
|
if output.next_token_logits is not None
|
|
else None
|
|
)
|
|
|
|
return LogitsProcessorOutput(
|
|
next_token_logits=next_token_logits,
|
|
full_logits=full_logits,
|
|
hidden_states=(
|
|
output.hidden_states[: self.raw_num_token]
|
|
if output.hidden_states is not None
|
|
else None
|
|
),
|
|
customized_info=output.customized_info,
|
|
)
|
|
else:
|
|
assert isinstance(output, PPProxyTensors)
|
|
return PPProxyTensors({k: v[: self.bs] for k, v in output.tensors.items()})
|
|
|
|
def get_spec_info(self, num_tokens: int):
|
|
spec_info = None
|
|
if (
|
|
self.model_runner.spec_algorithm.is_eagle()
|
|
or self.model_runner.spec_algorithm.is_standalone()
|
|
):
|
|
from sglang.srt.speculative.eagle_info import EagleVerifyInput
|
|
|
|
if self.model_runner.is_draft_worker:
|
|
raise RuntimeError("This should not happen.")
|
|
else:
|
|
|
|
capture_mode = (
|
|
CaptureHiddenMode.NULL
|
|
if self.model_runner.spec_algorithm.is_standalone()
|
|
else CaptureHiddenMode.FULL
|
|
)
|
|
spec_info = EagleVerifyInput(
|
|
draft_token=None,
|
|
custom_mask=self.buffers.custom_mask,
|
|
positions=None,
|
|
retrieve_index=None,
|
|
retrieve_next_token=None,
|
|
retrieve_next_sibling=None,
|
|
retrieve_cum_len=None,
|
|
spec_steps=self.speculative_num_steps,
|
|
topk=self.model_runner.server_args.speculative_eagle_topk,
|
|
draft_token_num=self.speculative_num_draft_tokens,
|
|
capture_hidden_mode=capture_mode,
|
|
seq_lens_sum=None,
|
|
seq_lens_cpu=None,
|
|
)
|
|
elif self.model_runner.spec_algorithm.is_dflash():
|
|
from sglang.srt.speculative.dflash_info import DFlashVerifyInput
|
|
from sglang.srt.speculative.dflash_utils import (
|
|
resolve_dflash_verify_mask_policy,
|
|
)
|
|
|
|
# Avoid enabling custom-mask modes during graph capture for backends that
|
|
# can express DFLASH verify via their built-in causal path.
|
|
_, build_custom_mask = resolve_dflash_verify_mask_policy(
|
|
self.model_runner.attn_backend
|
|
)
|
|
spec_info = DFlashVerifyInput(
|
|
draft_token=None,
|
|
positions=None,
|
|
draft_token_num=self.model_runner.server_args.speculative_num_draft_tokens,
|
|
custom_mask=(
|
|
None
|
|
if (self.model_runner.is_draft_worker or not build_custom_mask)
|
|
else self.buffers.custom_mask
|
|
),
|
|
capture_hidden_mode=(
|
|
CaptureHiddenMode.NULL
|
|
if self.model_runner.is_draft_worker
|
|
else CaptureHiddenMode.FULL
|
|
),
|
|
)
|
|
|
|
elif self.model_runner.spec_algorithm.is_ngram():
|
|
from sglang.srt.speculative.ngram_info import NgramVerifyInput
|
|
|
|
spec_info = NgramVerifyInput(
|
|
draft_token=None,
|
|
tree_mask=self.buffers.custom_mask,
|
|
positions=None,
|
|
retrieve_index=None,
|
|
retrieve_next_token=None,
|
|
retrieve_next_sibling=None,
|
|
draft_token_num=self.num_tokens_per_bs,
|
|
)
|
|
spec_info.capture_hidden_mode = CaptureHiddenMode.NULL
|
|
|
|
return spec_info
|
|
|
|
|
|
CUDA_GRAPH_CAPTURE_FAILED_MSG = (
|
|
"Possible solutions:\n"
|
|
"1. set --mem-fraction-static to a smaller value (e.g., 0.8 or 0.7)\n"
|
|
"2. set --cuda-graph-max-bs to a smaller value (e.g., 16)\n"
|
|
"3. disable torch compile by not using --enable-torch-compile\n"
|
|
"4. disable CUDA graph by --disable-cuda-graph. (Not recommended. Huge performance loss)\n"
|
|
"Open an issue on GitHub https://github.com/sgl-project/sglang/issues/new/choose \n"
|
|
)
|
|
|
|
|
|
class DeepEPCudaGraphRunnerAdapter:
|
|
def __init__(self):
|
|
# Record DeepEP mode used during capture to ensure replay consistency
|
|
self._captured_deepep_mode = None
|
|
|
|
def capture(self, is_extend_in_batch: bool):
|
|
if not get_moe_a2a_backend().is_deepep():
|
|
return
|
|
self._captured_deepep_mode = get_deepep_mode().resolve(
|
|
is_extend_in_batch=is_extend_in_batch
|
|
)
|
|
DeepEPBuffer.set_dispatch_mode(self._captured_deepep_mode)
|
|
|
|
def replay(self):
|
|
if not get_moe_a2a_backend().is_deepep():
|
|
return
|
|
assert self._captured_deepep_mode is not None
|
|
DeepEPBuffer.set_dispatch_mode(self._captured_deepep_mode)
|