minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Lifan Shen
2026-09-18 17:09:30 -07:00
committed by GitHub
parent b876213548
commit 8ea0ee300d
7 changed files with 395 additions and 158 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/sampler.py
+26 -18
View File
@@ -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."
)
python/sglang/srt/sampling/custom_logit_processor.py
+5 -1
View File
@@ -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
+80 -77
View File
@@ -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()