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:
laoyao0822
2026-06-27 20:37:05 +08:00
parent 132561a151
commit 2de822661c
2 changed files with 51 additions and 2 deletions

View File

@@ -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]