From 9df72e8f5a4bf03854aaf17f846eab1887ea139c Mon Sep 17 00:00:00 2001
From: jasonjk-park <jasonjk@fb.com>
Date: Thu, 10 Sep 2026 17:52:08 -0700
Subject: [PATCH] Fix custom logit processor params when num_tokens_in_batch is
 used (#38730)

---
 python/sglang/srt/layers/sampler.py           |  7 ++-
 .../sampling/test_custom_logit_processor.py   | 44 +++++++++++++++++++
 2 files changed, 50 insertions(+), 1 deletion(-)

diff --git a/python/sglang/srt/layers/sampler.py b/python/sglang/srt/layers/sampler.py
index 0b9323adf..d5ee8d826 100644
--- a/python/sglang/srt/layers/sampler.py
+++ b/python/sglang/srt/layers/sampler.py
@@ -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(
diff --git a/test/registered/unit/sampling/test_custom_logit_processor.py b/test/registered/unit/sampling/test_custom_logit_processor.py
index 843b36b55..df3e0dcad 100644
--- a/test/registered/unit/sampling/test_custom_logit_processor.py
+++ b/test/registered/unit/sampling/test_custom_logit_processor.py
@@ -12,6 +12,7 @@ from unittest.mock import MagicMock
 
 import torch
 
+from sglang.srt.layers.sampler import apply_custom_logit_processor
 from sglang.srt.parser.inkling_tokenizer import (
     CONTENT_THINKING,
     END_MESSAGE,
@@ -26,6 +27,7 @@ from sglang.srt.sampling.custom_logit_processor import (
     Qwen3ThinkingBudgetLogitProcessor,
     _cache_from_str,
 )
+from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
 from sglang.test.test_utils import CustomTestCase
 
 
@@ -37,6 +39,48 @@ def _make_req(origin_input_ids=None, output_ids=None):
     return req
 
 
+class TestApplyCustomLogitProcessor(CustomTestCase):
+    def test_repeats_request_params_for_each_token(self):
+        batch_size = 3
+        num_tokens = 2
+        params = [{"value": 1.0}, {"value": 2.0}, {"value": 3.0}]
+
+        def processor(logits, custom_params):
+            self.assertEqual(
+                custom_params, [params[0], params[0], params[2], params[2]]
+            )
+            for row, param in zip(logits, custom_params, strict=True):
+                row.fill_(param["value"])
+            return logits
+
+        sampling_info = SamplingBatchInfo(
+            temperatures=torch.ones(batch_size, 1),
+            top_ps=torch.ones(batch_size),
+            top_ks=torch.zeros(batch_size, dtype=torch.int32),
+            min_ps=torch.zeros(batch_size),
+            is_all_greedy=False,
+            is_any_greedy=False,
+            need_top_p_sampling=False,
+            need_top_k_sampling=False,
+            need_min_p_sampling=False,
+            vocab_size=4,
+            has_custom_logit_processor=True,
+            custom_params=params,
+            custom_logit_processor={0: (processor, torch.tensor([True, False, True]))},
+            device="cpu",
+        )
+        logits = torch.zeros(batch_size * num_tokens, 4)
+
+        apply_custom_logit_processor(
+            logits, sampling_info, num_tokens_in_batch=num_tokens
+        )
+
+        expected = torch.tensor(
+            [[1.0] * 4, [1.0] * 4, [0.0] * 4, [0.0] * 4, [3.0] * 4, [3.0] * 4]
+        )
+        self.assertTrue(torch.equal(logits, expected))
+
+
 # Serialization round-trip
 class TestCustomLogitProcessorSerialization(CustomTestCase):
     def test_to_str_produces_valid_json(self):