diff --git a/python/sglang/srt/layers/sampler.py b/python/sglang/srt/layers/sampler.py index abf59fd1c..197bf7dff 100644 --- a/python/sglang/srt/layers/sampler.py +++ b/python/sglang/srt/layers/sampler.py @@ -949,30 +949,38 @@ def apply_custom_logit_processor( f"({num_tokens_in_batch})" ) - for _, ( - processor, - batch_mask, - ) in sampling_batch_info.custom_logit_processor.items(): - # Get the batch indices that need to be processed - batch_indices = batch_mask.nonzero(as_tuple=True)[0] + batch_size = len(sampling_batch_info) + assert len(sampling_batch_info.custom_params) == batch_size, ( + f"The number of custom params ({len(sampling_batch_info.custom_params)}) does " + f"not match the number of sampling_batch_info ({batch_size})" + ) - assert batch_mask.shape[0] == len(sampling_batch_info), ( - f"The number of batch mask ({batch_mask.shape[0]}) does not match the number of " - f"sampling_batch_info ({len(sampling_batch_info)})" + token_offsets = ( + None + 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) ] - - # Apply the processor to the logits - logits[batch_mask] = processor( - logits[batch_mask], - custom_params, - ) + result = entry.processor(selected, custom_params) + logits.index_copy_(0, indices, result.to(logits.dtype)) logger.debug( - f"Custom logit processor {processor.__class__.__name__} is applied." + f"Custom logit processor {entry.processor.__class__.__name__} is applied." ) diff --git a/python/sglang/srt/sampling/custom_logit_processor.py b/python/sglang/srt/sampling/custom_logit_processor.py index 48ea39661..6c5a8c955 100644 --- a/python/sglang/srt/sampling/custom_logit_processor.py +++ b/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 diff --git a/python/sglang/srt/sampling/sampling_batch_info.py b/python/sglang/srt/sampling/sampling_batch_info.py index 966e91d4d..3d5c69de2 100644 --- a/python/sglang/srt/sampling/sampling_batch_info.py +++ b/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[ - Dict[int, Tuple[CustomLogitProcessor, torch.Tensor]] - ] = None + custom_logit_processor: Optional[Dict[int, ProcessorEntry]] = 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): ( - # The deserialized custom logit processor object - CustomLogitProcessor.from_str(processor_str), - # The mask tensor for the requests that use this custom logit processor - torch.zeros(len(reqs), dtype=torch.bool) - .scatter_(0, torch.tensor(true_indices), True) - .to(device, non_blocking=True), + hash(processor_str): ProcessorEntry( + processor=CustomLogitProcessor.from_str(processor_str), + rows=rows, + indices=_rows_to_device_indices(rows, device, _pin), ) - 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( - indices, - dtype=torch.long, - pin_memory=is_pin_memory_available(device), - ).to(device, non_blocking=True) + return _rows_to_device_indices(indices, device) 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( - self, keep_indices: List[int], keep_indices_device: torch.Tensor - ): + def _filter_batch_custom_logit_processor(self, keep_indices: List[int]): """Filter the custom logit processor and custom params""" - self.custom_logit_processor = { - k: (p, mask[keep_indices_device]) - for k, (p, mask) in self.custom_logit_processor.items() - if torch.any( - mask[keep_indices_device] - ) # ignore the custom logit processor whose mask is all False - } + position = {old: new for new, old in enumerate(keep_indices)} + pin = is_pin_memory_available(self.device) + kept = {} + for key, entry in self.custom_logit_processor.items(): + new_rows = sorted(position[old] for old in entry.rows if old in position) + 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 = ( - SamplingBatchInfo.merge_custom_logit_processor( - self.custom_logit_processor, - other.custom_logit_processor, - len(self), - len(other), - self.device, - ) - ) + self.custom_logit_processor = self._merge_processor_entries(other) # 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() diff --git a/test/registered/kernels/ops/sampling/test_apply_custom_logit_processor.py b/test/registered/kernels/ops/sampling/test_apply_custom_logit_processor.py new file mode 100644 index 000000000..4eab75fa5 --- /dev/null +++ b/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() diff --git a/test/registered/sampling/test_sampling_mask.py b/test/registered/sampling/test_sampling_mask.py index 187e805dc..f43faa1c8 100644 --- a/test/registered/sampling/test_sampling_mask.py +++ b/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: ( - DisallowedTokensLogitsProcessor(), - torch.tensor([True, False], device="cuda"), + 0: ProcessorEntry( + processor=DisallowedTokensLogitsProcessor(), + rows=[0], + indices=torch.tensor([0], device="cuda"), ) }, return_sampling_masks=[True, True], diff --git a/test/registered/unit/sampling/test_custom_logit_processor.py b/test/registered/unit/sampling/test_custom_logit_processor.py index df3e0dcad..dc3b1b42b 100644 --- a/test/registered/unit/sampling/test_custom_logit_processor.py +++ b/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) diff --git a/test/registered/unit/sampling/test_sampling_batch_info.py b/test/registered/unit/sampling/test_sampling_batch_info.py index 7e52f64ef..016f6d234 100644 --- a/test/registered/unit/sampling/test_sampling_batch_info.py +++ b/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): - """Test that merging two None processor dicts returns None.""" - result = SamplingBatchInfo.merge_custom_logit_processor( - None, None, 2, 3, DEVICE - ) - self.assertIsNone(result) + def test_merge_preserves_processors_and_offsets_rows(self): + proc_a, proc_b = MagicMock(), MagicMock() + for left_rows, right_rows, expected in ( + ({}, {}, {}), + ({1: [0]}, {1: [0, 2]}, {1: [0, 2, 4]}), + ({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): - """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 + left.merge_batch(right) - def test_disjoint_keys(self): - """Test that disjoint processor keys are merged with zero-filled padding.""" - proc_a = MagicMock() - proc_b = MagicMock() - lhs = {1: (proc_a, torch.tensor([True, False]))} - rhs = {2: (proc_b, torch.tensor([True]))} - result = SamplingBatchInfo.merge_custom_logit_processor(lhs, rhs, 2, 1, DEVICE) - # Key 1: lhs mask [True, False] + zero-filled rhs [False] - 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) + self.assertEqual(set(left.custom_logit_processor or {}), set(expected)) + for key, expected_rows in expected.items(): + entry = left.custom_logit_processor[key] + self.assertEqual(entry.rows, expected_rows) + self.assertEqual(entry.indices.tolist(), expected_rows) + if key not in right_rows: + self.assertIs(entry.indices, original_indices[key]) + self.assertIs(entry.processor, proc_a if key == 1 else proc_b) # 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] - self.assertEqual(mask.shape[0], 2) + entry = info.custom_logit_processor[42] + 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] - self.assertIsInstance(proc, DisallowedTokensLogitsProcessor) - self.assertTrue(mask[0].item()) - self.assertFalse(mask[1].item()) + entry = info.custom_logit_processor[key] + self.assertIsInstance(entry.processor, DisallowedTokensLogitsProcessor) + self.assertEqual(entry.rows, [0]) + 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)