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

Optimize conflicts between CUDA graph and vocab mask tensors (#1392)

Browse Source
This commit is contained in:
Liangsheng Yin
2024-09-13 20:27:53 -07:00
committed by GitHub
parent f3d32f888a
commit 70b6802982
32 changed files with 103 additions and 224 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -547,8 +547,9 @@ class ModelTpServer:
if self.model_runner.is_generation:
# Forward and sample the next tokens
if batch.extend_num_tokens != 0:
sample_output, logits_output = self.model_runner.forward(batch)
next_token_ids = batch.check_sample_results(sample_output)
logits_output = self.model_runner.forward(batch)
next_token_ids = self.model_runner.sample(logits_output, batch)
batch.sampling_info.penalizer_orchestrator.cumulate_output_tokens(
next_token_ids
)
@@ -723,8 +724,8 @@ class ModelTpServer:
batch.prepare_for_decode()
# Forward and sample the next tokens
sample_output, logits_output = self.model_runner.forward(batch)
next_token_ids = batch.check_sample_results(sample_output)
logits_output = self.model_runner.forward(batch)
next_token_ids = self.model_runner.sample(logits_output, batch)
batch.sampling_info.penalizer_orchestrator.cumulate_output_tokens(
next_token_ids
)