Fix custom logit processor params when num_tokens_in_batch is used (#38730)

2026-09-10 17:52:08 -07:00
parent d0035da34e
commit 9df72e8f5a
2 changed files with 50 additions and 1 deletions
@@ -961,11 +961,16 @@ def apply_custom_logit_processor(
            f"sampling_batch_info ({len(sampling_batch_info)})"
        )
        batch_mask = torch.repeat_interleave(batch_mask, num_tokens_in_batch)
+        custom_params = [
+            sampling_batch_info.custom_params[i]
+            for i in batch_indices
+            for _ in range(num_tokens_in_batch)
+        ]

        # Apply the processor to the logits
        logits[batch_mask] = processor(
            logits[batch_mask],
-            [sampling_batch_info.custom_params[i] for i in batch_indices],
+            custom_params,
        )

        logger.debug(