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

Remove unnecessary kernels of num_token_non_padded (#6965)

Browse Source
This commit is contained in:
fzyzcjy
2025-06-08 20:09:17 +08:00
committed by GitHub
parent 20d3ad3b58
commit 2fc1299562
3 changed files with 33 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -35,6 +35,7 @@ from sglang.srt.model_executor.forward_batch_info import (
ForwardBatch,
ForwardMode,
PPProxyTensors,
enable_num_token_non_padded,
)
from sglang.srt.patch_torch import monkey_patch_torch_compile
from sglang.srt.two_batch_overlap import TboCudaGraphRunnerPlugin
@@ -190,6 +191,9 @@ class CudaGraphRunner:
self.is_encoder_decoder = model_runner.model_config.is_encoder_decoder
self.enable_dp_attention = model_runner.server_args.enable_dp_attention
self.enable_sp_layernorm = model_runner.server_args.enable_sp_layernorm
self.enable_two_batch_overlap = (
model_runner.server_args.enable_two_batch_overlap
)
self.speculative_algorithm = model_runner.server_args.speculative_algorithm
self.tp_size = model_runner.server_args.tp_size
self.dp_size = model_runner.server_args.dp_size
@@ -327,9 +331,7 @@ class CudaGraphRunner:
)
is_tbo_supported = (
forward_batch.can_run_tbo
if self.model_runner.server_args.enable_two_batch_overlap
else True
forward_batch.can_run_tbo if self.enable_two_batch_overlap else True
)
return is_bs_supported and is_encoder_lens_supported and is_tbo_supported
@@ -549,13 +551,7 @@ class CudaGraphRunner:
self.seq_lens[:raw_bs].copy_(forward_batch.seq_lens)
self.out_cache_loc[:raw_num_token].copy_(forward_batch.out_cache_loc)
self.positions[:raw_num_token].copy_(forward_batch.positions)
num_token_non_padded = len(forward_batch.input_ids)
self.num_token_non_padded[...] = num_token_non_padded
self.tbo_plugin.replay_prepare(
forward_mode=forward_batch.forward_mode,
bs=bs,
num_token_non_padded=num_token_non_padded,
)
if forward_batch.seq_lens_cpu is not None:
if bs != raw_bs:
self.seq_lens_cpu.fill_(1)
@@ -572,6 +568,14 @@ class CudaGraphRunner:
self.mrope_positions[:, :raw_bs].copy_(forward_batch.mrope_positions)
if self.enable_dp_attention or self.enable_sp_layernorm:
self.global_num_tokens_gpu.copy_(forward_batch.global_num_tokens_gpu)
if enable_num_token_non_padded(self.model_runner.server_args):
self.num_token_non_padded.copy_(forward_batch.num_token_non_padded)
if self.enable_two_batch_overlap:
self.tbo_plugin.replay_prepare(
forward_mode=forward_batch.forward_mode,
bs=bs,
num_token_non_padded=len(forward_batch.input_ids),
)
# Attention backend
self.model_runner.attn_backend.init_forward_metadata_replay_cuda_graph(