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

Simplify pytorch sampling kernel and logit processor (#2491)

Browse Source
This commit is contained in:
Lianmin Zheng
2024-12-16 14:11:09 -08:00
committed by GitHub
parent 82699474fd
commit 7a1aecb938
5 changed files with 188 additions and 159 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -387,8 +387,14 @@ class CudaGraphRunner:
# Extract logprobs
if forward_batch.return_logprob:
next_token_logprobs = torch.nn.functional.log_softmax(
next_token_logits, dim=-1
logits_metadata = LogitsMetadata(
forward_mode=ForwardMode.DECODE,
top_logprobs_nums=forward_batch.top_logprobs_nums,
)
next_token_logprobs = (
LogitsProcessor.compute_temp_top_p_normalized_logprobs(
next_token_logits, logits_metadata
)
)
logits_output = LogitsProcessorOutput(
next_token_logits=next_token_logits,
@@ -396,10 +402,6 @@ class CudaGraphRunner:
)
return_top_logprob = any(x > 0 for x in forward_batch.top_logprobs_nums)
if return_top_logprob:
logits_metadata = LogitsMetadata(
forward_mode=ForwardMode.DECODE,
top_logprobs_nums=forward_batch.top_logprobs_nums,
)
(
logits_output.output_top_logprobs_val,
logits_output.output_top_logprobs_idx,