perf(sampling): avoid GPU syncs when applying custom logit processors (#39234)

python/sglang/srt/layers/sampler.py

@@ -949,30 +949,38 @@ def apply_custom_logit_processor(
         f"({num_tokens_in_batch})"
     )
 
-    for _, (
+    batch_size = len(sampling_batch_info)
-        processor,
+    assert len(sampling_batch_info.custom_params) == batch_size, (
-        batch_mask,
+        f"The number of custom params ({len(sampling_batch_info.custom_params)}) does "
-    ) in sampling_batch_info.custom_logit_processor.items():
+        f"not match the number of sampling_batch_info ({batch_size})"
-        # Get the batch indices that need to be processed
+    )
-        batch_indices = batch_mask.nonzero(as_tuple=True)[0]
 
-        assert batch_mask.shape[0] == len(sampling_batch_info), (
+    token_offsets = (
-            f"The number of batch mask ({batch_mask.shape[0]}) does not match the number of "
+        None
-            f"sampling_batch_info ({len(sampling_batch_info)})"
+        if num_tokens_in_batch == 1
+        else torch.arange(num_tokens_in_batch, device=sampling_batch_info.device)
+    )
+    for entry in sampling_batch_info.custom_logit_processor.values():
+        rows, indices = entry.rows, entry.indices
+        assert len(rows) == indices.numel(), (
+            f"The number of cached processor rows ({len(rows)}) does not match the "
+            f"number of cached device indices ({indices.numel()})"
         )
-        batch_mask = torch.repeat_interleave(batch_mask, num_tokens_in_batch)
+        assert not rows or rows[-1] < batch_size, (
+            f"Cached processor rows {rows} are stale for a batch of {batch_size}"
+        )
+
+        if token_offsets is not None:
+            indices = (indices[:, None] * num_tokens_in_batch + token_offsets).flatten()
+        selected = logits.index_select(0, indices)
         custom_params = [
             sampling_batch_info.custom_params[i]
-            for i in batch_indices
+            for i in rows
             for _ in range(num_tokens_in_batch)
         ]
-
+        result = entry.processor(selected, custom_params)
-        # Apply the processor to the logits
+        logits.index_copy_(0, indices, result.to(logits.dtype))
-        logits[batch_mask] = processor(
-            logits[batch_mask],
-            custom_params,
-        )
 
         logger.debug(
-            f"Custom logit processor {processor.__class__.__name__} is applied."
+            f"Custom logit processor {entry.processor.__class__.__name__} is applied."
         )

python/sglang/srt/sampling/custom_logit_processor.py

@@ -30,7 +30,11 @@ class CustomLogitProcessor(ABC):
         logits: torch.Tensor,
         custom_param_list: Optional[List[Dict[str, Any]]] = None,
     ) -> torch.Tensor:
-        """Define the callable behavior."""
+        """Define the callable behavior.
+
+        The returned tensor must have the same shape as `logits`: the caller
+        writes it back row for row and does not broadcast a reduced result.
+        """
         raise NotImplementedError
 
     @classmethod

python/sglang/srt/sampling/sampling_batch_info.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import dataclasses
 import logging
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
+from typing import TYPE_CHECKING, Any, Dict, List, Optional
 
 import torch
 
@@ -26,6 +26,26 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 
+def _rows_to_device_indices(
+    rows: List[int], device: str, pin: Optional[bool] = None
+) -> torch.Tensor:
+    """Move batch row numbers to the device without blocking on a pageable copy."""
+    if pin is None:
+        pin = is_pin_memory_available(device)
+    return torch.tensor(rows, dtype=torch.long, pin_memory=pin).to(
+        device, non_blocking=True
+    )
+
+
+@dataclasses.dataclass
+class ProcessorEntry:
+    """A custom logit processor and the batch rows it applies to."""
+
+    processor: CustomLogitProcessor
+    rows: List[int]
+    indices: torch.Tensor
+
+
 @dataclasses.dataclass
 class SamplingBatchInfo:
     # Basic batched sampling params
@@ -67,9 +87,7 @@ class SamplingBatchInfo:
     # Custom parameters
     custom_params: Optional[List[Optional[Dict[str, Any]]]] = None
     # Custom logit processor
-    custom_logit_processor: Optional[
+    custom_logit_processor: Optional[Dict[int, ProcessorEntry]] = None
-        Dict[int, Tuple[CustomLogitProcessor, torch.Tensor]]
-    ] = None
 
     # Used for deterministic sampling
     sampling_seed: Optional[torch.Tensor] = None
@@ -167,15 +185,12 @@ class SamplingBatchInfo:
                 processor_dict[processor_str].append(i)
 
             merged_custom_logit_processor = {
-                hash(processor_str): (
+                hash(processor_str): ProcessorEntry(
-                    # The deserialized custom logit processor object
+                    processor=CustomLogitProcessor.from_str(processor_str),
-                    CustomLogitProcessor.from_str(processor_str),
+                    rows=rows,
-                    # The mask tensor for the requests that use this custom logit processor
+                    indices=_rows_to_device_indices(rows, device, _pin),
-                    torch.zeros(len(reqs), dtype=torch.bool)
-                    .scatter_(0, torch.tensor(true_indices), True)
-                    .to(device, non_blocking=True),
                 )
-                for processor_str, true_indices in processor_dict.items()
+                for processor_str, rows in processor_dict.items()
             }
             custom_params = [r.sampling_params.custom_params for r in reqs]
         else:
@@ -251,11 +266,7 @@ class SamplingBatchInfo:
         ]
         if not indices:
             return None
-        return torch.tensor(
+        return _rows_to_device_indices(indices, device)
-            indices,
-            dtype=torch.long,
-            pin_memory=is_pin_memory_available(device),
-        ).to(device, non_blocking=True)
 
     def __len__(self):
         return len(self.temperatures)
@@ -348,7 +359,7 @@ class SamplingBatchInfo:
         self.penalizer_orchestrator.filter(keep_indices_device)
 
         if self.has_custom_logit_processor:
-            self._filter_batch_custom_logit_processor(keep_indices, keep_indices_device)
+            self._filter_batch_custom_logit_processor(keep_indices)
 
         for item in [
             "temperatures",
@@ -376,17 +387,25 @@ class SamplingBatchInfo:
 
         self.adjusted_filter_batch(keep_indices, keep_indices_device)
 
-    def _filter_batch_custom_logit_processor(
+    def _filter_batch_custom_logit_processor(self, keep_indices: List[int]):
-        self, keep_indices: List[int], keep_indices_device: torch.Tensor
-    ):
         """Filter the custom logit processor and custom params"""
-        self.custom_logit_processor = {
+        position = {old: new for new, old in enumerate(keep_indices)}
-            k: (p, mask[keep_indices_device])
+        pin = is_pin_memory_available(self.device)
-            for k, (p, mask) in self.custom_logit_processor.items()
+        kept = {}
-            if torch.any(
+        for key, entry in self.custom_logit_processor.items():
-                mask[keep_indices_device]
+            new_rows = sorted(position[old] for old in entry.rows if old in position)
-            )  # ignore the custom logit processor whose mask is all False
+            if not new_rows:
-        }
+                continue
+            kept[key] = (
+                entry
+                if new_rows == entry.rows
+                else ProcessorEntry(
+                    processor=entry.processor,
+                    rows=new_rows,
+                    indices=_rows_to_device_indices(new_rows, self.device, pin),
+                )
+            )
+        self.custom_logit_processor = kept
         self.custom_params = [self.custom_params[i] for i in keep_indices]
 
         # If the custom logit processor is an empty dict, set the flag to False,
@@ -396,61 +415,13 @@ class SamplingBatchInfo:
             self.custom_params = None
             self.has_custom_logit_processor = False
 
-    @staticmethod
-    def merge_custom_logit_processor(
-        lhs: Optional[Dict[int, Tuple[CustomLogitProcessor, torch.Tensor]]],
-        rhs: Optional[Dict[int, Tuple[CustomLogitProcessor, torch.Tensor]]],
-        bs1: int,
-        bs2: int,
-        device: str,
-    ):
-        if lhs is None and rhs is None:
-            return None
-        lhs, rhs = lhs or {}, rhs or {}
-
-        keys = set(lhs.keys()).union(set(rhs.keys()))
-        merged_dict = {}
-
-        for k in keys:
-            # Get the logit processor object
-            processor = lhs[k][0] if k in lhs else rhs[k][0]
-            # Get and merge the mask tensors from the two dicts
-            left_mask = (
-                lhs[k][1]
-                if k in lhs
-                else torch.zeros(bs1, dtype=torch.bool, device=device)
-            )
-            right_mask = (
-                rhs[k][1]
-                if k in rhs
-                else torch.zeros(bs2, dtype=torch.bool, device=device)
-            )
-            merged_dict[k] = (processor, torch.cat([left_mask, right_mask]))
-
-            assert merged_dict[k][1].shape[0] == bs1 + bs2, (
-                f"The batch size of merged mask ({merged_dict[k][1].shape[0]}) does not match "
-                f"the sum of the batch sizes of the two masks ({bs1 + bs2})"
-                f"\n{left_mask=}\n{right_mask=}\n{bs1=}\n{bs2=}"
-                f"\n{lhs=}\n{rhs=}"
-            )
-
-        return merged_dict
-
     def merge_batch(self, other: SamplingBatchInfo):
         self.penalizer_orchestrator.merge(other.penalizer_orchestrator)
 
         # Merge the custom logit processors and custom params lists
         if self.has_custom_logit_processor or other.has_custom_logit_processor:
             # Merge the custom logit processors
-            self.custom_logit_processor = (
+            self.custom_logit_processor = self._merge_processor_entries(other)
-                SamplingBatchInfo.merge_custom_logit_processor(
-                    self.custom_logit_processor,
-                    other.custom_logit_processor,
-                    len(self),
-                    len(other),
-                    self.device,
-                )
-            )
             # Merge the custom params lists
             self.custom_params = self.custom_params or [None] * len(self)
             other.custom_params = other.custom_params or [None] * len(other)
@@ -513,6 +484,38 @@ class SamplingBatchInfo:
 
         self.adjusted_merge_batch(other)
 
+    def _merge_processor_entries(
+        self, other: SamplingBatchInfo
+    ) -> Dict[int, ProcessorEntry]:
+        """Merge both batches' processor entries, shifting the right batch's rows."""
+        # This runs before temperatures are concatenated, so len(self) is the left batch size.
+        left = self.custom_logit_processor or {}
+        right = other.custom_logit_processor or {}
+        offset = len(self)
+
+        merged = {}
+        for key in left.keys() | right.keys():
+            left_entry = left.get(key)
+            right_entry = right.get(key)
+            if right_entry is None:
+                merged[key] = left_entry
+                continue
+            shifted_rows = [row + offset for row in right_entry.rows]
+            shifted_indices = right_entry.indices + offset
+            if left_entry is None:
+                merged[key] = ProcessorEntry(
+                    processor=right_entry.processor,
+                    rows=shifted_rows,
+                    indices=shifted_indices,
+                )
+            else:
+                merged[key] = ProcessorEntry(
+                    processor=left_entry.processor,
+                    rows=left_entry.rows + shifted_rows,
+                    indices=torch.cat([left_entry.indices, shifted_indices]),
+                )
+        return merged
+
     def copy_for_forward(self):
         # Accumulate the penalty into a pre-allocated buffer to get rid of the dependency of `penalizer_orchestrator` later
         self.update_penalties()

test/registered/kernels/ops/sampling/test_apply_custom_logit_processor.py

@@ -0,0 +1,113 @@
+import dataclasses
+import unittest
+
+import torch
+
+from sglang.srt.layers.sampler import apply_custom_logit_processor
+from sglang.srt.sampling.sampling_batch_info import ProcessorEntry, SamplingBatchInfo
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=15, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+
+class TestApplyCustomLogitProcessorCUDA(CustomTestCase):
+    def _make_info(self):
+        def processor(logits, params):
+            for row, param in zip(logits, params, strict=True):
+                row.narrow(0, param["token_id"], 1).fill_(-float("inf"))
+            return logits
+
+        return SamplingBatchInfo(
+            temperatures=torch.ones(3, 1, device="cuda"),
+            top_ps=torch.ones(3, device="cuda"),
+            top_ks=torch.zeros(3, dtype=torch.int32, device="cuda"),
+            min_ps=torch.zeros(3, device="cuda"),
+            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=[{"token_id": 1}, None, {"token_id": 2}],
+            custom_logit_processor={
+                0: ProcessorEntry(
+                    processor=processor,
+                    rows=[0, 2],
+                    indices=torch.tensor([0, 2], device="cuda"),
+                )
+            },
+            device="cuda",
+        )
+
+    @unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
+    def test_cached_decode_does_not_read_row_indices_back_to_cpu(self):
+        info = self._make_info()
+        for width in (1, 3):
+            with self.subTest(width=width):
+                logits = torch.zeros(3 * width, 4, device="cuda")
+                apply_custom_logit_processor(logits, info, width)
+                torch.cuda.synchronize()
+                with torch.profiler.profile(
+                    activities=[torch.profiler.ProfilerActivity.CPU]
+                ) as profile:
+                    apply_custom_logit_processor(logits, info, width)
+                names = {event.key for event in profile.key_averages()}
+                self.assertNotIn("aten::nonzero", names)
+                self.assertNotIn("aten::_local_scalar_dense", names)
+                expected = torch.zeros(3, 4)
+                expected[0, 1] = -float("inf")
+                expected[2, 2] = -float("inf")
+                self.assertTrue(
+                    torch.equal(logits.cpu(), expected.repeat_interleave(width, dim=0))
+                )
+
+    @unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
+    def test_casts_processor_result_to_logits_dtype(self):
+        def processor(logits, params):
+            return logits.float() + 0.1
+
+        info = self._make_info()
+        info.custom_logit_processor = {
+            0: dataclasses.replace(info.custom_logit_processor[0], processor=processor)
+        }
+        for width in (1, 3):
+            with self.subTest(width=width):
+                logits = torch.zeros(3 * width, 4, dtype=torch.bfloat16, device="cuda")
+                apply_custom_logit_processor(logits, info, width)
+                expected = torch.zeros(3, 4, dtype=torch.bfloat16)
+                expected[0] = 0.1
+                expected[2] = 0.1
+                self.assertEqual(logits.dtype, torch.bfloat16)
+                self.assertTrue(
+                    torch.equal(logits.cpu(), expected.repeat_interleave(width, dim=0))
+                )
+
+    @unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
+    def test_filter_does_not_read_processor_membership_back_to_cpu(self):
+        for keep in ([2, 1], [1]):
+            with self.subTest(keep=keep):
+                info = self._make_info()
+                self._make_info()._filter_batch_custom_logit_processor(keep)
+                torch.cuda.synchronize()
+                with torch.profiler.profile(
+                    activities=[torch.profiler.ProfilerActivity.CPU]
+                ) as profile:
+                    info._filter_batch_custom_logit_processor(keep)
+                names = {event.key for event in profile.key_averages()}
+                self.assertNotIn("aten::nonzero", names)
+                self.assertNotIn("aten::_local_scalar_dense", names)
+                if keep == [2, 1]:
+                    entry = info.custom_logit_processor[0]
+                    self.assertEqual(entry.rows, [0])
+                    self.assertEqual(entry.indices.tolist(), [0])
+                    self.assertEqual(info.custom_params, [{"token_id": 2}, None])
+                    self.assertEqual(set(info.custom_logit_processor), {0})
+                else:
+                    self.assertIsNone(info.custom_logit_processor)
+                    self.assertFalse(info.has_custom_logit_processor)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/sampling/test_sampling_mask.py

@@ -20,7 +20,10 @@ from sglang.srt.sampling.custom_logit_processor import (
     DisallowedTokensLogitsProcessor,
     Qwen3ThinkingBudgetLogitProcessor,
 )
-from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
+from sglang.srt.sampling.sampling_batch_info import (
+    ProcessorEntry,
+    SamplingBatchInfo,
+)
 from sglang.srt.utils import is_hip, kill_process_tree
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.test.test_utils import (
@@ -162,9 +165,10 @@ class TestSamplingMaskCapture(CustomTestCase):
                     has_custom_logit_processor=True,
                     custom_params=[{"token_ids": [2]}, None],
                     custom_logit_processor={
-                        0: (
+                        0: ProcessorEntry(
-                            DisallowedTokensLogitsProcessor(),
+                            processor=DisallowedTokensLogitsProcessor(),
-                            torch.tensor([True, False], device="cuda"),
+                            rows=[0],
+                            indices=torch.tensor([0], device="cuda"),
                         )
                     },
                     return_sampling_masks=[True, True],

test/registered/unit/sampling/test_custom_logit_processor.py

@@ -27,7 +27,10 @@ from sglang.srt.sampling.custom_logit_processor import (
     Qwen3ThinkingBudgetLogitProcessor,
     _cache_from_str,
 )
-from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
+from sglang.srt.sampling.sampling_batch_info import (
+    ProcessorEntry,
+    SamplingBatchInfo,
+)
 from sglang.test.test_utils import CustomTestCase
 
 
@@ -66,7 +69,11 @@ class TestApplyCustomLogitProcessor(CustomTestCase):
             vocab_size=4,
             has_custom_logit_processor=True,
             custom_params=params,
-            custom_logit_processor={0: (processor, torch.tensor([True, False, True]))},
+            custom_logit_processor={
+                0: ProcessorEntry(
+                    processor=processor, rows=[0, 2], indices=torch.tensor([0, 2])
+                )
+            },
             device="cpu",
         )
         logits = torch.zeros(batch_size * num_tokens, 4)

test/registered/unit/sampling/test_sampling_batch_info.py

@@ -12,7 +12,9 @@ from unittest.mock import MagicMock, patch
 import torch
 
 from sglang.srt.constrained.base_grammar_backend import GrammarMask
+from sglang.srt.sampling.custom_logit_processor import DisallowedTokensLogitsProcessor
 from sglang.srt.sampling.sampling_batch_info import (
+    ProcessorEntry,
     SamplingBatchInfo,
     merge_bias_tensor,
 )
@@ -149,48 +151,50 @@ class TestSamplingMaskBatchIndices(CustomTestCase):
 
 
 class TestMergeCustomLogitProcessor(CustomTestCase):
-    def test_both_none_returns_none(self):
+    def test_merge_preserves_processors_and_offsets_rows(self):
-        """Test that merging two None processor dicts returns None."""
+        proc_a, proc_b = MagicMock(), MagicMock()
-        result = SamplingBatchInfo.merge_custom_logit_processor(
+        for left_rows, right_rows, expected in (
-            None, None, 2, 3, DEVICE
+            ({}, {}, {}),
-        )
+            ({1: [0]}, {1: [0, 2]}, {1: [0, 2, 4]}),
-        self.assertIsNone(result)
+            ({1: [0]}, {2: [0]}, {1: [0], 2: [2]}),
+            ({}, {2: [0]}, {2: [2]}),
+            ({1: [1]}, {}, {1: [1]}),
+        ):
+            with self.subTest(left=left_rows, right=right_rows):
+                infos = []
+                for size, rows in ((2, left_rows), (3, right_rows)):
+                    infos.append(
+                        _make_info(
+                            batch_size=size,
+                            has_custom_logit_processor=bool(rows),
+                            custom_logit_processor={
+                                key: ProcessorEntry(
+                                    processor=proc_a if key == 1 else proc_b,
+                                    rows=values,
+                                    indices=torch.tensor(values, dtype=torch.long),
+                                )
+                                for key, values in rows.items()
+                            }
+                            or None,
+                            custom_params=[None] * size if rows else None,
+                        )
+                    )
+                left, right = infos
+                original_indices = {
+                    key: entry.indices
+                    for key, entry in (left.custom_logit_processor or {}).items()
+                }
 
-    def test_same_key_merges_masks(self):
+                left.merge_batch(right)
-        """Test that same processor key concatenates the boolean masks."""
-        proc = MagicMock()
-        lhs = {42: (proc, torch.tensor([True, False]))}
-        rhs = {42: (proc, torch.tensor([False, True, True]))}
-        result = SamplingBatchInfo.merge_custom_logit_processor(lhs, rhs, 2, 3, DEVICE)
-        self.assertIn(42, result)
-        self.assertEqual(result[42][1].shape[0], 5)
-        self.assertTrue(result[42][1][0].item())  # from lhs
-        self.assertFalse(result[42][1][1].item())  # from lhs
-        self.assertTrue(result[42][1][3].item())  # from rhs
 
-    def test_disjoint_keys(self):
+                self.assertEqual(set(left.custom_logit_processor or {}), set(expected))
-        """Test that disjoint processor keys are merged with zero-filled padding."""
+                for key, expected_rows in expected.items():
-        proc_a = MagicMock()
+                    entry = left.custom_logit_processor[key]
-        proc_b = MagicMock()
+                    self.assertEqual(entry.rows, expected_rows)
-        lhs = {1: (proc_a, torch.tensor([True, False]))}
+                    self.assertEqual(entry.indices.tolist(), expected_rows)
-        rhs = {2: (proc_b, torch.tensor([True]))}
+                    if key not in right_rows:
-        result = SamplingBatchInfo.merge_custom_logit_processor(lhs, rhs, 2, 1, DEVICE)
+                        self.assertIs(entry.indices, original_indices[key])
-        # Key 1: lhs mask [True, False] + zero-filled rhs [False]
+                    self.assertIs(entry.processor, proc_a if key == 1 else proc_b)
-        self.assertEqual(result[1][1].shape[0], 3)
-        self.assertTrue(result[1][1][0].item())
-        self.assertFalse(result[1][1][2].item())
-        # Key 2: zero-filled lhs [False, False] + rhs mask [True]
-        self.assertEqual(result[2][1].shape[0], 3)
-        self.assertFalse(result[2][1][0].item())
-        self.assertTrue(result[2][1][2].item())
-
-    def test_lhs_none_rhs_present(self):
-        """Test that None lhs is treated as empty dict and rhs mask is padded."""
-        proc = MagicMock()
-        rhs = {10: (proc, torch.tensor([True]))}
-        result = SamplingBatchInfo.merge_custom_logit_processor(None, rhs, 2, 1, DEVICE)
-        self.assertIn(10, result)
-        self.assertEqual(result[10][1].shape[0], 3)
 
 
 # apply_logits_bias
@@ -439,26 +443,100 @@ class TestFilterBatch(CustomTestCase):
         self.assertEqual(info.logit_bias.shape, (2, VOCAB_SIZE))
 
     def test_filter_with_custom_logit_processor(self):
-        """Test that filter updates both custom_params list and processor mask."""
         proc = MagicMock()
         info = _make_info(batch_size=3)
         info.has_custom_logit_processor = True
-        info.custom_logit_processor = {42: (proc, torch.tensor([True, False, True]))}
+        info.custom_logit_processor = {
+            42: ProcessorEntry(
+                processor=proc, rows=[0, 2], indices=torch.tensor([0, 2])
+            )
+        }
         info.custom_params = [{"a": 1}, {"b": 2}, {"c": 3}]
         keep = torch.tensor([0, 2])
         info.filter_batch([0, 2], keep)
         self.assertEqual(info.custom_params, [{"a": 1}, {"c": 3}])
-        mask = info.custom_logit_processor[42][1]
+        entry = info.custom_logit_processor[42]
-        self.assertEqual(mask.shape[0], 2)
+        self.assertEqual(entry.rows, [0, 1])
+        self.assertEqual(entry.indices.tolist(), entry.rows)
+
+    def test_filter_reuses_indices_only_when_rows_are_unchanged(self):
+        for keep, expected_rows in (([0, 1], [0, 1]), ([1, 2], [0]), ([2], [])):
+            with self.subTest(keep=keep):
+                original_indices = torch.tensor([0, 1])
+                info = _make_info(
+                    batch_size=3,
+                    has_custom_logit_processor=True,
+                    custom_logit_processor={
+                        42: ProcessorEntry(
+                            processor=MagicMock(),
+                            rows=[0, 1],
+                            indices=original_indices,
+                        )
+                    },
+                    custom_params=[None] * 3,
+                )
+                info.filter_batch(keep, torch.tensor(keep))
+                if not expected_rows:
+                    self.assertIsNone(info.custom_logit_processor)
+                    continue
+                entry = info.custom_logit_processor[42]
+                self.assertEqual(entry.rows, expected_rows)
+                self.assertEqual(entry.indices.tolist(), expected_rows)
+                if expected_rows == [0, 1]:
+                    self.assertIs(entry.indices, original_indices)
+                else:
+                    self.assertIsNot(entry.indices, original_indices)
+
+    def test_filter_merge_preserves_per_token_params(self):
+        from sglang.srt.layers.sampler import apply_custom_logit_processor
+
+        def processor(logits, params):
+            for row, param in zip(logits, params, strict=True):
+                row.fill_(param["value"])
+            return logits
+
+        info = _make_info(
+            batch_size=3,
+            has_custom_logit_processor=True,
+            custom_logit_processor={
+                42: ProcessorEntry(
+                    processor=processor, rows=[0, 2], indices=torch.tensor([0, 2])
+                )
+            },
+            custom_params=[{"value": 10}, None, {"value": 20}],
+        )
+        info.filter_batch([2, 1, 0], torch.tensor([2, 1, 0]))
+        info.merge_batch(
+            _make_info(
+                batch_size=1,
+                has_custom_logit_processor=True,
+                custom_logit_processor={
+                    42: ProcessorEntry(
+                        processor=processor, rows=[0], indices=torch.tensor([0])
+                    )
+                },
+                custom_params=[{"value": 30}],
+            )
+        )
+        entry = info.custom_logit_processor[42]
+        self.assertEqual(entry.rows, [0, 2, 3])
+        self.assertEqual(entry.indices.tolist(), entry.rows)
+        for width in (1, 3):
+            with self.subTest(width=width):
+                logits = torch.zeros(4 * width, VOCAB_SIZE)
+                apply_custom_logit_processor(logits, info, width)
+                expected = torch.tensor([20, 0, 10, 30]).repeat_interleave(width)
+                self.assertTrue(torch.equal(logits[:, 0], expected))
 
     def test_filter_removes_all_custom_processors(self):
         """Test cleanup when filter removes all requests using a processor."""
         proc = MagicMock()
         info = _make_info(batch_size=3)
         info.has_custom_logit_processor = True
-        info.custom_logit_processor = {42: (proc, torch.tensor([False, True, False]))}
+        info.custom_logit_processor = {
+            42: ProcessorEntry(processor=proc, rows=[1], indices=torch.tensor([1]))
+        }
         info.custom_params = [None, {"x": 1}, None]
-        # Keep only index 0 and 2 — processor 42's mask becomes [False, False]
         keep = torch.tensor([0, 2])
         info.filter_batch([0, 2], keep)
         self.assertFalse(info.has_custom_logit_processor)
@@ -522,7 +600,9 @@ class TestMergeBatch(CustomTestCase):
         proc = MagicMock()
         info1 = _make_info(batch_size=1)
         info1.has_custom_logit_processor = True
-        info1.custom_logit_processor = {1: (proc, torch.tensor([True]))}
+        info1.custom_logit_processor = {
+            1: ProcessorEntry(processor=proc, rows=[0], indices=torch.tensor([0]))
+        }
         info1.custom_params = [{"a": 1}]
         info2 = _make_info(batch_size=1)
         info2.has_custom_logit_processor = False
@@ -704,12 +784,30 @@ class TestFromScheduleBatch(CustomTestCase):
         info = SamplingBatchInfo.from_schedule_batch(batch, VOCAB_SIZE)
         self.assertIsNone(info.logit_bias)
 
+    def test_merge_preserves_processor_cache_after_batch_without_processors(self):
+        self._exec_ns.features.enable_custom_logit_processor = True
+        left_batch = MagicMock()
+        left_batch.reqs = [self._make_req(), self._make_req()]
+        left_batch.device = DEVICE
+        left = SamplingBatchInfo.from_schedule_batch(left_batch, VOCAB_SIZE)
+
+        processor_str = DisallowedTokensLogitsProcessor.to_str()
+        req = self._make_req()
+        req.custom_logit_processor = processor_str
+        req.sampling_params.custom_params = {"token_ids": [1]}
+        right_batch = MagicMock()
+        right_batch.reqs = [req]
+        right_batch.device = DEVICE
+        right = SamplingBatchInfo.from_schedule_batch(right_batch, VOCAB_SIZE)
+
+        left.merge_batch(right)
+
+        entry = left.custom_logit_processor[hash(processor_str)]
+        self.assertEqual(entry.rows, [2])
+        self.assertEqual(entry.indices.tolist(), [2])
+
     def test_custom_logit_processor_merging(self):
         """Test deserialization and merging of custom logit processors."""
-        from sglang.srt.sampling.custom_logit_processor import (
-            DisallowedTokensLogitsProcessor,
-        )
-
         self._exec_ns.features.enable_custom_logit_processor = True
 
         proc_str = DisallowedTokensLogitsProcessor.to_str()
@@ -728,12 +826,12 @@ class TestFromScheduleBatch(CustomTestCase):
         self.assertTrue(info.has_custom_logit_processor)
         self.assertIsNotNone(info.custom_logit_processor)
         self.assertEqual(len(info.custom_logit_processor), 1)
-        # Check the mask: req1 has processor (True), req2 doesn't (False)
         key = list(info.custom_logit_processor.keys())[0]
-        proc, mask = info.custom_logit_processor[key]
+        entry = info.custom_logit_processor[key]
-        self.assertIsInstance(proc, DisallowedTokensLogitsProcessor)
+        self.assertIsInstance(entry.processor, DisallowedTokensLogitsProcessor)
-        self.assertTrue(mask[0].item())
+        self.assertEqual(entry.rows, [0])
-        self.assertFalse(mask[1].item())
+        self.assertEqual(entry.indices.tolist(), entry.rows)
+        self.assertEqual(entry.indices.dtype, torch.long)
         # custom_params should be collected for all reqs
         self.assertEqual(len(info.custom_params), 2)