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

Speed up when having padding tokens in DeepEP (#6175)

Browse Source
This commit is contained in:
fzyzcjy
2025-05-18 07:44:05 +08:00
committed by GitHub
parent e3bed74afb
commit 2716830802
4 changed files with 53 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -240,6 +240,7 @@ class CudaGraphRunner:
self.out_cache_loc = torch.zeros((self.max_num_token,), dtype=torch.int64)
self.positions = torch.zeros((self.max_num_token,), dtype=torch.int64)
self.mrope_positions = torch.zeros((3, self.max_bs), dtype=torch.int64)
self.num_token_non_padded = torch.zeros((1,), dtype=torch.int32)
# pipeline parallelism
if self.pp_size > 1:
@@ -403,6 +404,7 @@ class CudaGraphRunner:
else:
encoder_lens = None
mrope_positions = self.mrope_positions[:, :bs]
self.num_token_non_padded[...] = num_tokens
# pipeline parallelism
if self.pp_size > 1:
@@ -461,6 +463,7 @@ class CudaGraphRunner:
spec_info=spec_info,
capture_hidden_mode=self.capture_hidden_mode,
lora_paths=lora_paths,
num_token_non_padded=self.num_token_non_padded,
)
if lora_paths is not None:
@@ -556,6 +559,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)
self.num_token_non_padded[...] = len(forward_batch.input_ids)
if forward_batch.seq_lens_cpu is not None:
if bs != raw_bs:
self.seq_lens_cpu.fill_(1)