Generalize auxiliary outputs (#39164)

Co-authored-by: jasonjk <jasonjk@twshared0085.36.lco2.facebook.com>

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,

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)

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__)
 

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
 
 

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, Optional, Protocol
-from typing import TYPE_CHECKING, Any, Callable, Optional, Protocol, Sequence
 
 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.
 

python/sglang/srt/sampling/sampling_observer_pp.py

@@ -11,10 +11,8 @@ from typing import (
 
 import torch
 
-from sglang.srt.sampling.sampling_observer import (
+from sglang.srt.managers.auxiliary_output import DeviceAuxiliaryOutput
-    DeviceAuxiliaryOutput,
+from sglang.srt.sampling.sampling_observer import SamplingObserver
-    SamplingObserver,
-)
 
 
 @runtime_checkable

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(