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

[Perf] Avoid per-decode-step host sync in min_new_tokens penalty (#28397)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-06-15 23:32:28 -07:00
committed by GitHub
parent 637c9f780b
commit 556cf54d47
1 changed files with 5 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/sampling/penaltylib/min_new_tokens.py
+5 -2
View File
@@ -67,8 +67,11 @@ class BatchedMinNewTokensPenalizer(_BatchedPenalizer):
self.len_output_tokens += 1 self.len_output_tokens += 1
def _apply(self, logits: torch.Tensor): def _apply(self, logits: torch.Tensor):
mask = (self.len_output_tokens < self.min_new_tokens).expand_as(logits) # Boolean-mask indexing (logits[mask]) is data-dependent and forces a
logits[mask] += self.stop_token_penalties[mask] # device-to-host sync every decode step; torch.where is a plain
# elementwise select with no sync (and no -inf*0=nan).
mask = self.len_output_tokens < self.min_new_tokens
logits.add_(torch.where(mask, self.stop_token_penalties, 0.0))
def _filter(self, keep_indices: torch.Tensor): def _filter(self, keep_indices: torch.Tensor):
self.min_new_tokens = self.min_new_tokens[keep_indices] self.min_new_tokens = self.min_new_tokens[keep_indices]