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

Log if cuda graph is used & extend cuda graph capture to cuda-graph-max-bs (#6201)

Browse Source 
Co-authored-by: SangBin Cho <rkooo567@gmail.com>
This commit is contained in:
Lianmin Zheng
2025-05-12 00:17:33 -07:00
committed by GitHub
parent 7d3a3d4510
commit fba8eccd7e
27 changed files with 293 additions and 121 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
python/sglang/srt/constrained/base_grammar_backend.py
View File

@@ -37,6 +37,12 @@ class BaseGrammarObject:
"""
raise NotImplementedError()
def rollback(self, k: int):
raise NotImplementedError()
def is_terminated(self):
raise NotImplementedError()
def allocate_vocab_mask(
self, vocab_size: int, batch_size: int, device
) -> torch.Tensor:

4
python/sglang/srt/disaggregation/prefill.py
View File

@@ -277,19 +277,17 @@ class SchedulerDisaggregationPrefillMixin:
next_token_ids,
extend_input_len_per_req,
extend_logprob_start_len_per_req,
bid,
) = (
result.logits_output,
result.next_token_ids,
result.extend_input_len_per_req,
result.extend_logprob_start_len_per_req,
result.bid,
)
# Transfer kv for prefill completed requests and add it into disagg_prefill_infight_queue
if self.enable_overlap:
# wait
_, next_token_ids = self.tp_worker.resolve_last_batch_result(launch_done)
_, next_token_ids, _ = self.tp_worker.resolve_last_batch_result(launch_done)
else:
next_token_ids = result.next_token_ids.tolist()

2
python/sglang/srt/entrypoints/engine.py
View File

@@ -330,7 +330,7 @@ class Engine(EngineBase):
return {
**dataclasses.asdict(self.tokenizer_manager.server_args),
**self.scheduler_info,
**internal_states,
"internal_states": internal_states,
"version": __version__,
}

2
python/sglang/srt/entrypoints/http_server.py
View File

@@ -222,7 +222,7 @@ async def get_server_info():
return {
**dataclasses.asdict(_global_state.tokenizer_manager.server_args),
**_global_state.scheduler_info,
**internal_states,
"internal_states": internal_states,
"version": __version__,
}

6
python/sglang/srt/layers/attention/utils.py
View File

@@ -28,7 +28,8 @@ def create_flashinfer_kv_indices_triton(
num_loop = tl.cdiv(kv_end - kv_start, BLOCK_SIZE)
for i in range(num_loop):
offset = tl.arange(0, BLOCK_SIZE) + i * BLOCK_SIZE
# index into req_to_token_ptr needs to be int64
offset = tl.arange(0, BLOCK_SIZE).to(tl.int64) + i * BLOCK_SIZE
mask = offset < kv_end - kv_start
data = tl.load(
req_to_token_ptr
@@ -70,8 +71,9 @@ def create_flashmla_kv_indices_triton(
num_pages_loop = tl.cdiv(kv_end - kv_start, BLOCK_SIZE)
for i in range(num_pages_loop):
# index into req_to_token_ptr needs to be int64
paged_offset = (
tl.arange(0, NUM_PAGE_PER_BLOCK) + i * NUM_PAGE_PER_BLOCK
tl.arange(0, NUM_PAGE_PER_BLOCK).to(tl.int64) + i * NUM_PAGE_PER_BLOCK
) * PAGED_SIZE
paged_offset_out = tl.arange(0, NUM_PAGE_PER_BLOCK) + i * NUM_PAGE_PER_BLOCK

28
python/sglang/srt/managers/scheduler.py
View File

@@ -160,6 +160,7 @@ class GenerationBatchResult:
extend_input_len_per_req: List[int]
extend_logprob_start_len_per_req: List[int]
bid: int
can_run_cuda_graph: bool
@dataclass
@@ -323,13 +324,14 @@ class Scheduler(
set_random_seed(self.random_seed)
# Print debug info
logger.info(
f"max_total_num_tokens={self.max_total_num_tokens}, "
f"chunked_prefill_size={server_args.chunked_prefill_size}, "
f"max_prefill_tokens={self.max_prefill_tokens}, "
f"max_running_requests={self.max_running_requests}, "
f"context_len={self.model_config.context_len}"
)
if tp_rank == 0:
logger.info(
f"max_total_num_tokens={self.max_total_num_tokens}, "
f"chunked_prefill_size={server_args.chunked_prefill_size}, "
f"max_prefill_tokens={self.max_prefill_tokens}, "
f"max_running_requests={self.max_running_requests}, "
f"context_len={self.model_config.context_len}"
)
# Init memory pool and cache
self.init_memory_pool_and_cache()
@@ -752,6 +754,7 @@ class Scheduler(
extend_input_len_per_req=None,
extend_logprob_start_len_per_req=None,
bid=bids[next_mb_id],
can_run_cuda_graph=result.can_run_cuda_graph,
)
self.process_batch_result(mbs[next_mb_id], output_result)
last_mbs[next_mb_id] = mbs[next_mb_id]
@@ -1159,7 +1162,9 @@ class Scheduler(
self.metrics_collector.log_stats(self.stats)
def log_decode_stats(self, running_batch=None):
def log_decode_stats(
self, can_run_cuda_graph: bool, running_batch: ScheduleBatch = None
):
batch = running_batch or self.running_batch
gap_latency = time.time() - self.last_decode_stats_tic
@@ -1199,6 +1204,7 @@ class Scheduler(
msg += f"pre-allocated usage: {self.num_tokens_pre_allocated / self.max_total_num_tokens:.2f}, "
msg += (
f"cuda graph: {can_run_cuda_graph}, "
f"gen throughput (token/s): {self.last_gen_throughput:.2f}, "
f"#queue-req: {len(self.waiting_queue)}"
)
@@ -1524,11 +1530,11 @@ class Scheduler(
if self.spec_algorithm.is_none():
model_worker_batch = batch.get_model_worker_batch()
if self.pp_group.is_last_rank:
logits_output, next_token_ids = (
logits_output, next_token_ids, can_run_cuda_graph = (
self.tp_worker.forward_batch_generation(model_worker_batch)
)
else:
pp_hidden_states_proxy_tensors, _ = (
pp_hidden_states_proxy_tensors, _, can_run_cuda_graph = (
self.tp_worker.forward_batch_generation(model_worker_batch)
)
bid = model_worker_batch.bid
@@ -1538,6 +1544,7 @@ class Scheduler(
next_token_ids,
bid,
num_accepted_tokens,
can_run_cuda_graph,
) = self.draft_worker.forward_batch_speculative_generation(batch)
self.spec_num_total_accepted_tokens += (
num_accepted_tokens + batch.batch_size()
@@ -1571,6 +1578,7 @@ class Scheduler(
extend_input_len_per_req=extend_input_len_per_req,
extend_logprob_start_len_per_req=extend_logprob_start_len_per_req,
bid=bid,
can_run_cuda_graph=can_run_cuda_graph,
)
else: # embedding or reward model
model_worker_batch = batch.get_model_worker_batch()

18
python/sglang/srt/managers/scheduler_output_processor_mixin.py
View File

@@ -38,20 +38,16 @@ class SchedulerOutputProcessorMixin:
next_token_ids,
extend_input_len_per_req,
extend_logprob_start_len_per_req,
bid,
) = (
result.logits_output,
result.next_token_ids,
result.extend_input_len_per_req,
result.extend_logprob_start_len_per_req,
result.bid,
)
if self.enable_overlap:
logits_output, next_token_ids = (
self.tp_worker.resolve_last_batch_result(
launch_done,
)
logits_output, next_token_ids, _ = (
self.tp_worker.resolve_last_batch_result(launch_done)
)
else:
# Move next_token_ids and logprobs to cpu
@@ -189,16 +185,16 @@ class SchedulerOutputProcessorMixin:
result: GenerationBatchResult,
launch_done: Optional[threading.Event] = None,
):
logits_output, next_token_ids, bid = (
logits_output, next_token_ids, can_run_cuda_graph = (
result.logits_output,
result.next_token_ids,
result.bid,
result.can_run_cuda_graph,
)
self.num_generated_tokens += len(batch.reqs)
if self.enable_overlap:
logits_output, next_token_ids = self.tp_worker.resolve_last_batch_result(
launch_done
logits_output, next_token_ids, can_run_cuda_graph = (
self.tp_worker.resolve_last_batch_result(launch_done)
)
next_token_logprobs = logits_output.next_token_logprobs
elif batch.spec_algorithm.is_none():
@@ -280,7 +276,7 @@ class SchedulerOutputProcessorMixin:
self.attn_tp_rank == 0
and self.forward_ct_decode % self.server_args.decode_log_interval == 0
):
self.log_decode_stats(running_batch=batch)
self.log_decode_stats(can_run_cuda_graph, running_batch=batch)
def add_input_logprob_return_values(
self: Scheduler,

7
python/sglang/srt/managers/tokenizer_manager.py
View File

@@ -923,12 +923,13 @@ class TokenizerManager:
):
await self.send_to_scheduler.send_pyobj(obj)
async def get_internal_state(self) -> Dict[Any, Any]:
async def get_internal_state(self) -> List[Dict[Any, Any]]:
req = GetInternalStateReq()
res: List[GetInternalStateReqOutput] = (
responses: List[GetInternalStateReqOutput] = (
await self.get_internal_state_communicator(req)
)
return res[0].internal_state
# Many DP ranks
return [res.internal_state for res in responses]
def get_log_request_metadata(self):
max_length = None

21
python/sglang/srt/managers/tp_worker.py
View File

@@ -20,7 +20,7 @@ from typing import Optional, Tuple, Union
import torch
from sglang.srt.configs.model_config import ModelConfig
from sglang.srt.distributed import get_pp_group, get_tp_group, get_world_group
from sglang.srt.distributed import get_pp_group, get_world_group
from sglang.srt.hf_transformers_utils import (
get_processor,
get_tokenizer,
@@ -183,8 +183,11 @@ class TpModelWorker:
def forward_batch_generation(
self,
model_worker_batch: ModelWorkerBatch,
launch_done: Optional[threading.Event] = None,
skip_sample: bool = False,
) -> Tuple[Union[LogitsProcessorOutput, torch.Tensor], Optional[torch.Tensor]]:
) -> Tuple[
Union[LogitsProcessorOutput, torch.Tensor], Optional[torch.Tensor], bool
]:
forward_batch = ForwardBatch.init_new(model_worker_batch, self.model_runner)
pp_proxy_tensors = None
@@ -196,11 +199,11 @@ class TpModelWorker:
)
if self.pp_group.is_last_rank:
logits_output = self.model_runner.forward(
logits_output, can_run_cuda_graph = self.model_runner.forward(
forward_batch, pp_proxy_tensors=pp_proxy_tensors
)
if model_worker_batch.launch_done is not None:
model_worker_batch.launch_done.set()
if launch_done is not None:
launch_done.set()
if skip_sample:
next_token_ids = None
@@ -209,17 +212,17 @@ class TpModelWorker:
logits_output, model_worker_batch
)
return logits_output, next_token_ids
return logits_output, next_token_ids, can_run_cuda_graph
else:
pp_proxy_tensors = self.model_runner.forward(
pp_proxy_tensors, can_run_cuda_graph = self.model_runner.forward(
forward_batch,
pp_proxy_tensors=pp_proxy_tensors,
)
return pp_proxy_tensors.tensors, None
return pp_proxy_tensors.tensors, None, can_run_cuda_graph
def forward_batch_embedding(self, model_worker_batch: ModelWorkerBatch):
forward_batch = ForwardBatch.init_new(model_worker_batch, self.model_runner)
logits_output = self.model_runner.forward(forward_batch)
logits_output, _ = self.model_runner.forward(forward_batch)
embeddings = logits_output.embeddings
return embeddings

24
python/sglang/srt/managers/tp_worker_overlap_thread.py
View File

@@ -18,7 +18,7 @@ import logging
import signal
import threading
from queue import Queue
from typing import Optional
from typing import Optional, Tuple
import psutil
import torch
@@ -145,8 +145,10 @@ class TpModelWorkerClient:
resolve_future_token_ids(input_ids, self.future_token_ids_map)
# Run forward
logits_output, next_token_ids = self.worker.forward_batch_generation(
model_worker_batch
logits_output, next_token_ids, can_run_cuda_graph = (
self.worker.forward_batch_generation(
model_worker_batch, model_worker_batch.launch_done
)
)
# Update the future token ids map
@@ -171,14 +173,18 @@ class TpModelWorkerClient:
next_token_ids = next_token_ids.to("cpu", non_blocking=True)
copy_done.record()
self.output_queue.put((copy_done, logits_output, next_token_ids))
self.output_queue.put(
(copy_done, logits_output, next_token_ids, can_run_cuda_graph)
)
def resolve_last_batch_result(self, launch_done: Optional[threading.Event] = None):
"""
This function is called to resolve the last batch result and
wait for the current batch to be launched. Used in overlap mode.
"""
copy_done, logits_output, next_token_ids = self.output_queue.get()
copy_done, logits_output, next_token_ids, can_run_cuda_graph = (
self.output_queue.get()
)
if launch_done is not None:
launch_done.wait()
@@ -193,9 +199,11 @@ class TpModelWorkerClient:
logits_output.input_token_logprobs.tolist()
)
next_token_ids = next_token_ids.tolist()
return logits_output, next_token_ids
return logits_output, next_token_ids, can_run_cuda_graph
def forward_batch_generation(self, model_worker_batch: ModelWorkerBatch):
def forward_batch_generation(
self, model_worker_batch: ModelWorkerBatch
) -> Tuple[None, torch.Tensor, bool]:
# Create a new copy of sampling_info because it will be updated in-place by the scheduler for the next batch.
sampling_info = model_worker_batch.sampling_info
sampling_info.update_penalties()
@@ -223,7 +231,7 @@ class TpModelWorkerClient:
self.future_token_ids_ct = (
self.future_token_ids_ct + bs
) % self.future_token_ids_limit
return None, future_next_token_ids
return None, future_next_token_ids, False
def update_weights_from_disk(self, recv_req: UpdateWeightFromDiskReqInput):
success, message = self.worker.update_weights_from_disk(recv_req)

17
python/sglang/srt/model_executor/cuda_graph_runner.py
View File

@@ -19,7 +19,7 @@ import bisect
import inspect
import os
from contextlib import contextmanager
from typing import TYPE_CHECKING, Callable
from typing import TYPE_CHECKING, Callable, Optional, Union
import torch
import tqdm
@@ -40,15 +40,12 @@ from sglang.srt.patch_torch import monkey_patch_torch_compile
from sglang.srt.utils import (
get_available_gpu_memory,
get_device_memory_capacity,
is_hip,
rank0_log,
)
if TYPE_CHECKING:
from sglang.srt.model_executor.model_runner import ModelRunner
_is_hip = is_hip()
def _to_torch(model: torch.nn.Module, reverse: bool, num_tokens: int):
for sub in model._modules.values():
@@ -137,7 +134,6 @@ def get_batch_sizes_to_capture(model_runner: ModelRunner):
)
gpu_mem = get_device_memory_capacity()
# Batch size of each rank will not become so large when DP is on
if gpu_mem is not None and gpu_mem > 96 * 1024:
capture_bs += list(range(160, 257, 8))
@@ -148,12 +144,15 @@ def get_batch_sizes_to_capture(model_runner: ModelRunner):
model_runner.req_to_token_pool.size
]
capture_bs = list(sorted(set(capture_bs)))
assert len(capture_bs) > 0 and capture_bs[0] > 0
capture_bs = [bs for bs in capture_bs if bs <= model_runner.req_to_token_pool.size]
if server_args.cuda_graph_max_bs:
capture_bs = [bs for bs in capture_bs if bs <= server_args.cuda_graph_max_bs]
if max(capture_bs) < server_args.cuda_graph_max_bs:
capture_bs += list(
range(max(capture_bs), server_args.cuda_graph_max_bs + 1, 16)
)
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
compile_bs = (
[bs for bs in capture_bs if bs <= server_args.torch_compile_max_bs]
if server_args.enable_torch_compile

15
python/sglang/srt/model_executor/model_runner.py
View File

@@ -1085,32 +1085,33 @@ class ModelRunner:
forward_batch: ForwardBatch,
skip_attn_backend_init: bool = False,
pp_proxy_tensors: Optional[PPProxyTensors] = None,
) -> Union[LogitsProcessorOutput, PPProxyTensors]:
) -> Tuple[Union[LogitsProcessorOutput, PPProxyTensors], bool]:
can_run_cuda_graph = bool(
forward_batch.forward_mode.is_cuda_graph()
and self.cuda_graph_runner
and self.cuda_graph_runner.can_run(forward_batch)
)
if can_run_cuda_graph:
return self.cuda_graph_runner.replay(
ret = self.cuda_graph_runner.replay(
forward_batch,
skip_attn_backend_init=skip_attn_backend_init,
pp_proxy_tensors=pp_proxy_tensors,
)
if forward_batch.forward_mode.is_decode():
return self.forward_decode(forward_batch, pp_proxy_tensors=pp_proxy_tensors)
elif forward_batch.forward_mode.is_decode():
ret = self.forward_decode(forward_batch, pp_proxy_tensors=pp_proxy_tensors)
elif forward_batch.forward_mode.is_extend():
return self.forward_extend(
ret = self.forward_extend(
forward_batch,
skip_attn_backend_init=skip_attn_backend_init,
pp_proxy_tensors=pp_proxy_tensors,
)
elif forward_batch.forward_mode.is_idle():
return self.forward_idle(forward_batch, pp_proxy_tensors=pp_proxy_tensors)
ret = self.forward_idle(forward_batch, pp_proxy_tensors=pp_proxy_tensors)
else:
raise ValueError(f"Invalid forward mode: {forward_batch.forward_mode}")
return ret, can_run_cuda_graph
def _preprocess_logits(
self, logits_output: LogitsProcessorOutput, sampling_info: SamplingBatchInfo
):

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

@@ -1086,7 +1086,7 @@ class ServerArgs:
"--cuda-graph-max-bs",
type=int,
default=ServerArgs.cuda_graph_max_bs,
help="Set the maximum batch size for cuda graph.",
help="Set the maximum batch size for cuda graph. It will extend the cuda graph capture batch size to this value.",
)
parser.add_argument(
"--cuda-graph-bs",

27
python/sglang/srt/speculative/eagle_worker.py
View File

@@ -251,8 +251,8 @@ class EAGLEWorker(TpModelWorker):
if batch.forward_mode.is_decode():
with self.draft_tp_context(self.draft_model_runner.tp_group):
spec_info = self.draft(batch)
logits_output, verify_output, model_worker_batch = self.verify(
batch, spec_info
logits_output, verify_output, model_worker_batch, can_run_cuda_graph = (
self.verify(batch, spec_info)
)
# If it is None, it means all requests are finished
@@ -264,21 +264,22 @@ class EAGLEWorker(TpModelWorker):
verify_output.verified_id,
model_worker_batch.bid,
sum(verify_output.accept_length_per_req_cpu),
can_run_cuda_graph,
)
elif batch.forward_mode.is_idle():
model_worker_batch = batch.get_model_worker_batch()
logits_output, next_token_ids = self.target_worker.forward_batch_generation(
model_worker_batch
logits_output, next_token_ids, _ = (
self.target_worker.forward_batch_generation(model_worker_batch)
)
return logits_output, next_token_ids, model_worker_batch.bid, 0
return logits_output, next_token_ids, model_worker_batch.bid, 0, False
else:
logits_output, next_token_ids, bid = self.forward_target_extend(batch)
with self.draft_tp_context(self.draft_model_runner.tp_group):
self.forward_draft_extend(
batch, logits_output.hidden_states, next_token_ids
)
return logits_output, next_token_ids, bid, 0
return logits_output, next_token_ids, bid, 0, False
def forward_target_extend(
self, batch: ScheduleBatch
@@ -297,7 +298,7 @@ class EAGLEWorker(TpModelWorker):
# We need the full hidden states to prefill the KV cache of the draft model.
model_worker_batch = batch.get_model_worker_batch()
model_worker_batch.capture_hidden_mode = CaptureHiddenMode.FULL
logits_output, next_token_ids = self.target_worker.forward_batch_generation(
logits_output, next_token_ids, _ = self.target_worker.forward_batch_generation(
model_worker_batch
)
return logits_output, next_token_ids, model_worker_batch.bid
@@ -478,8 +479,10 @@ class EAGLEWorker(TpModelWorker):
batch.forward_mode = ForwardMode.TARGET_VERIFY
batch.spec_info = spec_info
model_worker_batch = batch.get_model_worker_batch()
logits_output, _ = self.target_worker.forward_batch_generation(
model_worker_batch, skip_sample=True
logits_output, _, can_run_cuda_graph = (
self.target_worker.forward_batch_generation(
model_worker_batch, skip_sample=True
)
)
self._detect_nan_if_needed(logits_output)
spec_info.hidden_states = logits_output.hidden_states
@@ -504,7 +507,7 @@ class EAGLEWorker(TpModelWorker):
if batch.return_logprob:
self.add_logprob_values(batch, res, logits_output)
return logits_output, res, model_worker_batch
return logits_output, res, model_worker_batch, can_run_cuda_graph
def add_logprob_values(
self,
@@ -590,7 +593,7 @@ class EAGLEWorker(TpModelWorker):
model_worker_batch, self.draft_model_runner
)
forward_batch.return_logprob = False
logits_output = self.draft_model_runner.forward(forward_batch)
logits_output, _ = self.draft_model_runner.forward(forward_batch)
self._detect_nan_if_needed(logits_output)
assert isinstance(forward_batch.spec_info, EagleDraftInput)
assert forward_batch.spec_info is batch.spec_info
@@ -617,7 +620,7 @@ class EAGLEWorker(TpModelWorker):
)
# Run
logits_output = self.draft_model_runner.forward(forward_batch)
logits_output, _ = self.draft_model_runner.forward(forward_batch)
self._detect_nan_if_needed(logits_output)
self.capture_for_decode(logits_output, forward_batch.spec_info)