diff --git a/python/sglang/srt/layers/logits_processor.py b/python/sglang/srt/layers/logits_processor.py index b9835654a..6aac6eb21 100644 --- a/python/sglang/srt/layers/logits_processor.py +++ b/python/sglang/srt/layers/logits_processor.py @@ -51,13 +51,13 @@ from sglang.srt.layers.logprob_processor import ( get_top_logprobs_raw, ) from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding +from sglang.srt.managers.auxiliary_output import DeviceAuxiliaryOutput from sglang.srt.model_executor.forward_batch_info import ( CaptureHiddenMode, ForwardBatch, ForwardMode, ) from sglang.srt.runtime_context import get_exec, get_parallel -from sglang.srt.sampling.sampling_observer import DeviceAuxiliaryOutput from sglang.srt.utils.common import ( is_cpu, is_npu, diff --git a/python/sglang/srt/managers/auxiliary_output.py b/python/sglang/srt/managers/auxiliary_output.py new file mode 100644 index 000000000..31c1fa5f6 --- /dev/null +++ b/python/sglang/srt/managers/auxiliary_output.py @@ -0,0 +1,93 @@ +"""Auxiliary outputs carried through the generation-result lifecycle.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Callable, Optional, Protocol, Sequence + +import torch + +if TYPE_CHECKING: + from sglang.srt.managers.schedule_batch import ScheduleBatch + + +@dataclass(frozen=True) +class CommittedTokens: + output_index: int + token_ids: tuple[int, ...] + + +class DeviceAuxiliaryOutput(Protocol): + """Device output copied later by the scheduler. + + Tensors must not alias CUDA-graph static buffers that a later replay can + overwrite before the scheduler-side copy completes. + """ + + def copy_to_host( + self, copy_tensor: Callable[[torch.Tensor], torch.Tensor] + ) -> HostAuxiliaryOutput: ... + + +class HostAuxiliaryOutput(Protocol): + """Scheduler-side result produced by ``DeviceAuxiliaryOutput``. + + ``consume`` runs after sampled tokens have been committed to each request + and immediately before response streaming. ``commits`` is aligned with + ``batch.reqs`` and identifies only the newly visible tokens. Implementations + can buffer per-request values for a ``SchedulerOutputStreamer`` subclass to + expose through customized response metadata. + """ + + def consume( + self, + batch: ScheduleBatch, + commits: Sequence[Optional[CommittedTokens]], + ) -> None: ... + + +@dataclass(frozen=True) +class CompositeHostAuxiliaryOutput: + outputs: tuple[HostAuxiliaryOutput, ...] + + def consume( + self, + batch: ScheduleBatch, + commits: Sequence[Optional[CommittedTokens]], + ) -> None: + for output in self.outputs: + output.consume(batch, commits) + + +@dataclass(frozen=True) +class CompositeDeviceAuxiliaryOutput: + outputs: tuple[DeviceAuxiliaryOutput, ...] + + def copy_to_host( + self, copy_tensor: Callable[[torch.Tensor], torch.Tensor] + ) -> CompositeHostAuxiliaryOutput: + return CompositeHostAuxiliaryOutput( + tuple(output.copy_to_host(copy_tensor) for output in self.outputs) + ) + + +def append_auxiliary_output( + current: Optional[DeviceAuxiliaryOutput], + output: Optional[DeviceAuxiliaryOutput], +) -> Optional[DeviceAuxiliaryOutput]: + if output is None: + return current + if current is None: + return output + + current_outputs = ( + current.outputs + if isinstance(current, CompositeDeviceAuxiliaryOutput) + else (current,) + ) + new_outputs = ( + output.outputs + if isinstance(output, CompositeDeviceAuxiliaryOutput) + else (output,) + ) + return CompositeDeviceAuxiliaryOutput(current_outputs + new_outputs) diff --git a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py index 9fa4e9b85..386e8d10f 100644 --- a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py +++ b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py @@ -20,6 +20,7 @@ from sglang.srt.layers.logits_processor import ( LogitsProcessorOutput, SamplingMaskStatus, ) +from sglang.srt.managers.auxiliary_output import CommittedTokens from sglang.srt.managers.schedule_batch import ( FINISH_ABORT, FINISH_MATCHED_TOKEN, @@ -44,7 +45,6 @@ from sglang.srt.runtime_context import ( mamba_track_grid, max_speculative_num_draft_tokens, ) -from sglang.srt.sampling.sampling_observer import CommittedTokens from sglang.srt.sampling.sampling_params import ( get_request_reasoning_end_token_ids, ) @@ -58,6 +58,7 @@ if TYPE_CHECKING: from sglang.srt.disaggregation.decode_kvcache_offload_manager import ( DecodeKVCacheOffloadManager, ) + from sglang.srt.managers.auxiliary_output import HostAuxiliaryOutput from sglang.srt.managers.hisparse_coordinator import HiSparseCoordinator from sglang.srt.managers.scheduler_components.logprob_result_processor import ( SchedulerLogprobResultProcessor, @@ -77,7 +78,6 @@ if TYPE_CHECKING: from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache from sglang.srt.mem_cache.memory_pool import ReqToTokenPool from sglang.srt.observability.metrics_collector import SchedulerMetricsCollector - from sglang.srt.sampling.sampling_observer import HostAuxiliaryOutput logger = logging.getLogger(__name__) diff --git a/python/sglang/srt/managers/utils.py b/python/sglang/srt/managers/utils.py index 88ee08f5d..be44849ba 100644 --- a/python/sglang/srt/managers/utils.py +++ b/python/sglang/srt/managers/utils.py @@ -23,8 +23,8 @@ from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.srt.state_capturer.base import TopkCaptureOutput if TYPE_CHECKING: + from sglang.srt.managers.auxiliary_output import HostAuxiliaryOutput from sglang.srt.managers.scheduler import GenerationBatchResult - from sglang.srt.sampling.sampling_observer import HostAuxiliaryOutput from sglang.srt.speculative.spec_info import SpecInput diff --git a/python/sglang/srt/sampling/sampling_observer.py b/python/sglang/srt/sampling/sampling_observer.py index 135d196ee..94ef0a8d8 100644 --- a/python/sglang/srt/sampling/sampling_observer.py +++ b/python/sglang/srt/sampling/sampling_observer.py @@ -9,51 +9,16 @@ manager consumes the resulting customized response fields. from __future__ import annotations -from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Callable, Optional, Protocol, Sequence +from typing import TYPE_CHECKING, Any, Optional, Protocol import torch +from sglang.srt.managers.auxiliary_output import DeviceAuxiliaryOutput + if TYPE_CHECKING: - from sglang.srt.managers.schedule_batch import ScheduleBatch from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo -@dataclass(frozen=True) -class CommittedTokens: - output_index: int - token_ids: tuple[int, ...] - - -class DeviceAuxiliaryOutput(Protocol): - """Device output copied later by the scheduler. - - Tensors must not alias CUDA-graph static buffers that a later replay can - overwrite before the scheduler-side copy completes. - """ - - def copy_to_host( - self, copy_tensor: Callable[[torch.Tensor], torch.Tensor] - ) -> HostAuxiliaryOutput: ... - - -class HostAuxiliaryOutput(Protocol): - """Scheduler-side result produced by ``DeviceAuxiliaryOutput``. - - ``consume`` runs after sampled tokens have been committed to each request - and immediately before response streaming. ``commits`` is aligned with - ``batch.reqs`` and identifies only the newly visible tokens. Implementations - can buffer per-request values for a ``SchedulerOutputStreamer`` subclass to - expose through customized response metadata. - """ - - def consume( - self, - batch: ScheduleBatch, - commits: Sequence[Optional[CommittedTokens]], - ) -> None: ... - - class SamplingObserver(Protocol): """Invocation-scoped hooks around the production grammar mask and sampler. diff --git a/python/sglang/srt/sampling/sampling_observer_pp.py b/python/sglang/srt/sampling/sampling_observer_pp.py index 4de76b969..559a8f22e 100644 --- a/python/sglang/srt/sampling/sampling_observer_pp.py +++ b/python/sglang/srt/sampling/sampling_observer_pp.py @@ -11,10 +11,8 @@ from typing import ( import torch -from sglang.srt.sampling.sampling_observer import ( - DeviceAuxiliaryOutput, - SamplingObserver, -) +from sglang.srt.managers.auxiliary_output import DeviceAuxiliaryOutput +from sglang.srt.sampling.sampling_observer import SamplingObserver @runtime_checkable diff --git a/test/registered/unit/managers/test_generation_auxiliary_output.py b/test/registered/unit/managers/test_auxiliary_output.py similarity index 95% rename from test/registered/unit/managers/test_generation_auxiliary_output.py rename to test/registered/unit/managers/test_auxiliary_output.py index 21e6d449f..2afa825d2 100644 --- a/test/registered/unit/managers/test_generation_auxiliary_output.py +++ b/test/registered/unit/managers/test_auxiliary_output.py @@ -11,6 +11,10 @@ from sglang.srt.layers.logits_processor import ( SamplingMaskOutput, SamplingMaskStatus, ) +from sglang.srt.managers.auxiliary_output import ( + CompositeDeviceAuxiliaryOutput, + append_auxiliary_output, +) from sglang.srt.managers.scheduler import Scheduler from sglang.srt.managers.scheduler_components.batch_result_processor import ( SchedulerBatchResultProcessor, @@ -30,6 +34,10 @@ register_cpu_ci(est_time=14, suite="base-a-test-cpu") @dataclass class HostOutput: values: torch.Tensor + consumed: object = None + + def consume(self, batch, commits): + self.consumed = (batch, commits) class DeviceOutput: @@ -209,6 +217,40 @@ def test_auxiliary_host_outputs_are_owned_by_each_generation_result(): assert first_device.copy_count == second_device.copy_count == 1 +def test_composite_auxiliary_output_copies_and_consumes_each_child(): + first = DeviceOutput(torch.tensor([1.0])) + second = DeviceOutput(torch.tensor([2.0])) + third = DeviceOutput(torch.tensor([3.0])) + + device_output = append_auxiliary_output(first, second) + device_output = append_auxiliary_output(device_output, third) + + assert isinstance(device_output, CompositeDeviceAuxiliaryOutput) + assert device_output.outputs == (first, second, third) + + copied = [] + + def copy_tensor(tensor): + copied.append(tensor) + return tensor.clone() + + host_output = device_output.copy_to_host(copy_tensor) + batch = object() + commits = [object()] + host_output.consume(batch, commits) + + assert copied == [first.values, second.values, third.values] + assert first.copy_count == second.copy_count == third.copy_count == 1 + assert all(output.consumed == (batch, commits) for output in host_output.outputs) + + +def test_append_auxiliary_output_ignores_absent_outputs(): + output = DeviceOutput(torch.tensor([1.0])) + + assert append_auxiliary_output(None, output) is output + assert append_auxiliary_output(output, None) is output + + def test_sampling_clears_stale_device_output_when_observer_produces_no_state(): runner = _model_runner_for_sampling_path() runner.sampling_observer = SimpleNamespace(