Avoid decode-step syncs in min-new-token penalties
The old min_new_tokens penalizer updated logits through boolean-mask indexing. That indexing is data-dependent and can force synchronization on the decode hot path. Use an elementwise torch.where followed by inplace add so the operation stays tensorized and avoids the mask-index update path. Constraint: Keep the numeric behavior for active rows equivalent without multiplying by zero, which would turn -inf penalties into NaN. Rejected: Expanding the mask and assigning through logits[mask] | this is the synchronization pattern being removed. Confidence: high Scope-risk: narrow Directive: Do not reintroduce boolean-mask writes in decode-step penalty paths without profiling the synchronization behavior. Tested: RED/GREEN local pytest test/registered/unit/sampling/test_min_new_tokens_penalizer.py Tested: RED/GREEN remote pytest in cjy-glm5-new for min_new_tokens penalizer together with related regression tests Tested: git diff --check; py_compile for min_new_tokens.py Not-tested: Full decode throughput benchmark
This commit is contained in:
@@ -67,8 +67,8 @@ class BatchedMinNewTokensPenalizer(_BatchedPenalizer):
|
||||
self.len_output_tokens += 1
|
||||
|
||||
def _apply(self, logits: torch.Tensor):
|
||||
mask = (self.len_output_tokens < self.min_new_tokens).expand_as(logits)
|
||||
logits[mask] += self.stop_token_penalties[mask]
|
||||
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):
|
||||
self.min_new_tokens = self.min_new_tokens[keep_indices]
|
||||
|
||||
Reference in New Issue
Block a user