Mohammad Miadh AngkadandMohammad Angkad
22f02cc339
Co-authored-by: Mohammad Angkad <mohammad.angkad@radixark.ai>
903 lines
30 KiB
Python
903 lines
30 KiB
Python
from dataclasses import dataclass
|
|
from types import SimpleNamespace
|
|
from unittest.mock import Mock, patch
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from sglang.srt.disaggregation.prefill import SchedulerDisaggregationPrefillMixin
|
|
from sglang.srt.layers.logits_processor import (
|
|
LogitsProcessorOutput,
|
|
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,
|
|
)
|
|
from sglang.srt.managers.scheduler_pp_mixin import PPBatchMetadata
|
|
from sglang.srt.managers.utils import GenerationBatchResult, get_logprob_from_pp_outputs
|
|
from sglang.srt.model_executor.forward_batch_info import PPProxyTensors
|
|
from sglang.srt.model_executor.model_runner import ModelRunner
|
|
from sglang.srt.runtime_context import publish, reset_context
|
|
from sglang.srt.server_args import ServerArgs
|
|
from sglang.srt.speculative.eagle_info import EagleDraftInput
|
|
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
|
|
from sglang.test.ci.ci_register import register_cpu_ci
|
|
from sglang.test.test_utils import published_topology
|
|
|
|
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:
|
|
def __init__(self, values: torch.Tensor):
|
|
self.values = values
|
|
self.copy_count = 0
|
|
|
|
def copy_to_host(self, copy_tensor):
|
|
self.copy_count += 1
|
|
return HostOutput(copy_tensor(self.values))
|
|
|
|
def to_pp_tensors(self):
|
|
return {"values": self.values}
|
|
|
|
|
|
class HostOnlyDeviceOutput:
|
|
def __init__(self, values: torch.Tensor):
|
|
self.values = values
|
|
|
|
def copy_to_host(self, copy_tensor):
|
|
return HostOutput(copy_tensor(self.values))
|
|
|
|
|
|
class Observer:
|
|
def __init__(self):
|
|
self.received_tensors = None
|
|
|
|
def from_pp_tensors(self, tensors):
|
|
self.received_tensors = tensors
|
|
return DeviceOutput(tensors["values"])
|
|
|
|
|
|
class CopyDone:
|
|
def __init__(self):
|
|
self.record_count = 0
|
|
|
|
def record(self):
|
|
self.record_count += 1
|
|
|
|
|
|
def _model_runner_for_sampling_path(
|
|
*,
|
|
spec_algorithm=SpeculativeAlgorithm.NONE,
|
|
dllm_algorithm=None,
|
|
):
|
|
# `supports_sampling_observer` reads the dLLM algorithm from the bags, so
|
|
# the path is stated by publishing it rather than by standing one in.
|
|
reset_context()
|
|
publish(ServerArgs(model_path="dummy", dllm_algorithm=dllm_algorithm), role="test")
|
|
runner = object.__new__(ModelRunner)
|
|
runner.server_args = SimpleNamespace()
|
|
runner.spec_algorithm = spec_algorithm
|
|
runner._sampling_observer = None
|
|
return runner
|
|
|
|
|
|
def setup_function(_):
|
|
# The code under test reads its config from the bags.
|
|
reset_context()
|
|
publish(ServerArgs(model_path="dummy"), role="test")
|
|
|
|
|
|
def teardown_function(_):
|
|
reset_context()
|
|
|
|
|
|
def test_auxiliary_output_releases_device_holder_after_copy():
|
|
device_output = DeviceOutput(torch.tensor([1.0, 2.0]))
|
|
logits_output = LogitsProcessorOutput(
|
|
next_token_logits=None,
|
|
auxiliary_device_output=device_output,
|
|
)
|
|
result = GenerationBatchResult(
|
|
logits_output=logits_output,
|
|
next_token_ids=torch.tensor([7]),
|
|
copy_done=CopyDone(),
|
|
)
|
|
|
|
result.copy_to_cpu(return_logprob=False)
|
|
|
|
assert logits_output.auxiliary_device_output is None
|
|
assert result.auxiliary_host_output is not device_output
|
|
assert device_output.copy_count == 1
|
|
assert result.copy_done.record_count == 1
|
|
|
|
|
|
def test_sampling_mask_output_uses_generation_result_copy_path():
|
|
sampling_output = SamplingMaskOutput(
|
|
token_ids=torch.tensor([[3, 5]], dtype=torch.int32),
|
|
lengths=torch.tensor([2], dtype=torch.int32),
|
|
selected_logprobs=torch.tensor([-0.5]),
|
|
statuses=torch.tensor([SamplingMaskStatus.OK], dtype=torch.int32),
|
|
)
|
|
result = GenerationBatchResult(
|
|
logits_output=LogitsProcessorOutput(
|
|
next_token_logits=None,
|
|
sampling_mask_output=sampling_output,
|
|
),
|
|
next_token_ids=torch.tensor([3]),
|
|
copy_done=CopyDone(),
|
|
)
|
|
|
|
with patch(
|
|
"sglang.srt.managers.utils._async_d2h",
|
|
side_effect=lambda tensor: tensor.clone(),
|
|
) as copy_tensor:
|
|
result.copy_to_cpu(return_logprob=False)
|
|
|
|
assert copy_tensor.call_count == 5
|
|
assert sampling_output.token_ids.tolist() == [[3, 5]]
|
|
assert sampling_output.lengths.tolist() == [2]
|
|
assert sampling_output.statuses.tolist() == [SamplingMaskStatus.OK]
|
|
assert result.copy_done.record_count == 1
|
|
|
|
|
|
def test_pipeline_sampling_mask_round_trip_without_logprobs():
|
|
sampling_output = SamplingMaskOutput(
|
|
token_ids=torch.tensor([[3, 5]], dtype=torch.int32),
|
|
lengths=torch.tensor([2], dtype=torch.int32),
|
|
selected_logprobs=torch.tensor([-0.5]),
|
|
statuses=torch.tensor([SamplingMaskStatus.OK], dtype=torch.int32),
|
|
)
|
|
result = GenerationBatchResult(
|
|
logits_output=LogitsProcessorOutput(
|
|
next_token_logits=None, sampling_mask_output=sampling_output
|
|
),
|
|
next_token_ids=torch.tensor([3]),
|
|
)
|
|
payload = Scheduler._pp_prepare_tensor_dict(
|
|
object.__new__(Scheduler), result, SimpleNamespace(return_logprob=False)
|
|
)
|
|
output, _, _ = get_logprob_from_pp_outputs(PPProxyTensors(payload))
|
|
for name in ("token_ids", "lengths", "selected_logprobs", "statuses"):
|
|
torch.testing.assert_close(
|
|
getattr(output.sampling_mask_output, name), getattr(sampling_output, name)
|
|
)
|
|
|
|
|
|
def test_non_pp_auxiliary_output_only_requires_host_copy_support():
|
|
device_output = HostOnlyDeviceOutput(torch.tensor([1.0, 2.0]))
|
|
result = GenerationBatchResult(
|
|
logits_output=LogitsProcessorOutput(
|
|
next_token_logits=None,
|
|
auxiliary_device_output=device_output,
|
|
),
|
|
next_token_ids=torch.tensor([7]),
|
|
copy_done=CopyDone(),
|
|
)
|
|
|
|
result.copy_to_cpu(return_logprob=False)
|
|
|
|
assert torch.equal(result.auxiliary_host_output.values, device_output.values)
|
|
|
|
|
|
def test_auxiliary_host_outputs_are_owned_by_each_generation_result():
|
|
logits_output = LogitsProcessorOutput(next_token_logits=None)
|
|
first_device = DeviceOutput(torch.tensor([1.0]))
|
|
logits_output.auxiliary_device_output = first_device
|
|
first = GenerationBatchResult(
|
|
logits_output=logits_output,
|
|
next_token_ids=torch.tensor([1]),
|
|
copy_done=CopyDone(),
|
|
)
|
|
first.copy_to_cpu(return_logprob=False)
|
|
|
|
second_device = DeviceOutput(torch.tensor([2.0]))
|
|
logits_output.auxiliary_device_output = second_device
|
|
second = GenerationBatchResult(
|
|
logits_output=logits_output,
|
|
next_token_ids=torch.tensor([2]),
|
|
copy_done=CopyDone(),
|
|
)
|
|
second.copy_to_cpu(return_logprob=False)
|
|
|
|
assert first.auxiliary_host_output.values.tolist() == [1.0]
|
|
assert second.auxiliary_host_output.values.tolist() == [2.0]
|
|
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(
|
|
is_active=lambda sampling_info: True,
|
|
after_sample=lambda state, token_ids: None,
|
|
)
|
|
runner._preprocess_logits = Mock(return_value=None)
|
|
runner.sampler = lambda *args, **kwargs: torch.tensor([3])
|
|
runner.ngram_embedding_manager = SimpleNamespace(
|
|
update_after_decode=lambda **kwargs: None
|
|
)
|
|
logits_output = LogitsProcessorOutput(
|
|
next_token_logits=torch.zeros(1, 4),
|
|
auxiliary_device_output=DeviceOutput(torch.tensor([99.0])),
|
|
)
|
|
forward_batch = SimpleNamespace(
|
|
sampling_info=object(),
|
|
return_logprob=False,
|
|
top_logprobs_nums=None,
|
|
token_ids_logprobs=None,
|
|
positions=torch.tensor([0]),
|
|
seq_lens=torch.tensor([1]),
|
|
forward_mode=SimpleNamespace(is_decode=lambda: True),
|
|
)
|
|
|
|
ModelRunner.sample(runner, logits_output, forward_batch)
|
|
|
|
assert logits_output.auxiliary_device_output is None
|
|
runner._preprocess_logits.assert_called_once_with(
|
|
logits_output,
|
|
forward_batch.sampling_info,
|
|
observer=runner.sampling_observer,
|
|
)
|
|
|
|
|
|
def test_sampling_publishes_observer_output_for_the_sampled_tokens():
|
|
state = object()
|
|
device_output = DeviceOutput(torch.tensor([4.0]))
|
|
observer = SimpleNamespace(
|
|
is_active=Mock(return_value=True),
|
|
after_sample=Mock(return_value=device_output),
|
|
)
|
|
runner = _model_runner_for_sampling_path()
|
|
runner.sampling_observer = observer
|
|
runner._preprocess_logits = Mock(return_value=state)
|
|
sampled_tokens = torch.tensor([3])
|
|
runner.sampler = Mock(return_value=sampled_tokens)
|
|
runner.ngram_embedding_manager = SimpleNamespace(
|
|
update_after_decode=lambda **kwargs: None
|
|
)
|
|
logits_output = LogitsProcessorOutput(next_token_logits=torch.zeros(1, 4))
|
|
forward_batch = SimpleNamespace(
|
|
sampling_info=object(),
|
|
return_logprob=False,
|
|
top_logprobs_nums=None,
|
|
token_ids_logprobs=None,
|
|
positions=torch.tensor([0]),
|
|
seq_lens=torch.tensor([1]),
|
|
forward_mode=SimpleNamespace(is_decode=lambda: True),
|
|
)
|
|
|
|
ModelRunner.sample(runner, logits_output, forward_batch)
|
|
|
|
assert logits_output.auxiliary_device_output is device_output
|
|
runner._preprocess_logits.assert_called_once_with(
|
|
logits_output,
|
|
forward_batch.sampling_info,
|
|
observer=observer,
|
|
)
|
|
observer.is_active.assert_called_once_with(forward_batch.sampling_info)
|
|
observer.after_sample.assert_called_once_with(state, sampled_tokens)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("spec_algorithm", "dllm_algorithm"),
|
|
[
|
|
(SpeculativeAlgorithm.EAGLE, None),
|
|
(SpeculativeAlgorithm.NONE, "dream"),
|
|
],
|
|
)
|
|
def test_sampling_observer_rejects_sampling_paths_that_bypass_hooks(
|
|
spec_algorithm,
|
|
dllm_algorithm,
|
|
):
|
|
runner = _model_runner_for_sampling_path(
|
|
spec_algorithm=spec_algorithm,
|
|
dllm_algorithm=dllm_algorithm,
|
|
)
|
|
|
|
with pytest.raises(ValueError, match="configured sampling path"):
|
|
runner.sampling_observer = Observer()
|
|
|
|
|
|
def test_custom_sampling_path_can_enable_sampling_observer():
|
|
class SupportedModelRunner(ModelRunner):
|
|
def supports_sampling_observer(self):
|
|
return True
|
|
|
|
runner = object.__new__(SupportedModelRunner)
|
|
observer = Observer()
|
|
|
|
runner.sampling_observer = observer
|
|
|
|
assert runner.sampling_observer is observer
|
|
|
|
|
|
@pytest.mark.parametrize("has_inactive_observer", [False, True])
|
|
def test_sampling_without_active_observer_preserves_preprocess_override(
|
|
has_inactive_observer,
|
|
):
|
|
observer = (
|
|
SimpleNamespace(is_active=Mock(return_value=False))
|
|
if has_inactive_observer
|
|
else None
|
|
)
|
|
runner = _model_runner_for_sampling_path()
|
|
runner.sampling_observer = observer
|
|
runner._preprocess_logits = Mock(
|
|
side_effect=lambda logits_output, sampling_info: None
|
|
)
|
|
runner.sampler = Mock(return_value=torch.tensor([3]))
|
|
runner.ngram_embedding_manager = SimpleNamespace(
|
|
update_after_decode=lambda **kwargs: None
|
|
)
|
|
logits_output = LogitsProcessorOutput(next_token_logits=torch.zeros(1, 4))
|
|
forward_batch = SimpleNamespace(
|
|
sampling_info=object(),
|
|
return_logprob=False,
|
|
top_logprobs_nums=None,
|
|
token_ids_logprobs=None,
|
|
positions=torch.tensor([0]),
|
|
seq_lens=torch.tensor([1]),
|
|
forward_mode=SimpleNamespace(is_decode=lambda: True),
|
|
)
|
|
|
|
ModelRunner.sample(runner, logits_output, forward_batch)
|
|
|
|
runner._preprocess_logits.assert_called_once_with(
|
|
logits_output, forward_batch.sampling_info
|
|
)
|
|
if observer is not None:
|
|
observer.is_active.assert_called_once_with(forward_batch.sampling_info)
|
|
|
|
|
|
def test_preprocess_logits_without_observer_uses_standard_path():
|
|
runner = object.__new__(ModelRunner)
|
|
logits_output = LogitsProcessorOutput(next_token_logits=torch.zeros(1, 4))
|
|
grammar_mask = object()
|
|
sampling_info = SimpleNamespace(
|
|
grammar_mask=grammar_mask,
|
|
update_regex_vocab_mask=Mock(),
|
|
apply_logits_bias=Mock(),
|
|
apply_logits_bias_with_observer=Mock(),
|
|
)
|
|
|
|
state = ModelRunner._preprocess_logits(
|
|
runner,
|
|
logits_output,
|
|
sampling_info,
|
|
)
|
|
|
|
assert state is None
|
|
sampling_info.update_regex_vocab_mask.assert_called_once_with()
|
|
sampling_info.apply_logits_bias.assert_called_once_with(
|
|
logits_output.next_token_logits
|
|
)
|
|
sampling_info.apply_logits_bias_with_observer.assert_not_called()
|
|
assert sampling_info.grammar_mask is None
|
|
|
|
|
|
def test_active_observer_uses_observer_logits_preprocessing():
|
|
runner = object.__new__(ModelRunner)
|
|
observer = SimpleNamespace()
|
|
observer_state = object()
|
|
logits_output = LogitsProcessorOutput(next_token_logits=torch.zeros(1, 4))
|
|
sampling_info = SimpleNamespace(
|
|
grammar_mask=object(),
|
|
update_regex_vocab_mask=Mock(),
|
|
apply_logits_bias=Mock(),
|
|
apply_logits_bias_with_observer=Mock(return_value=observer_state),
|
|
)
|
|
|
|
state = ModelRunner._preprocess_logits(
|
|
runner,
|
|
logits_output,
|
|
sampling_info,
|
|
observer=observer,
|
|
)
|
|
|
|
assert state is observer_state
|
|
sampling_info.update_regex_vocab_mask.assert_called_once_with()
|
|
sampling_info.apply_logits_bias.assert_not_called()
|
|
sampling_info.apply_logits_bias_with_observer.assert_called_once_with(
|
|
logits_output.next_token_logits,
|
|
observer=observer,
|
|
)
|
|
assert sampling_info.grammar_mask is None
|
|
|
|
|
|
def test_scheduler_copies_auxiliary_output_for_non_overlap_results():
|
|
event = object()
|
|
scheduler = object.__new__(Scheduler)
|
|
scheduler.device_module = SimpleNamespace(Event=Mock(return_value=event))
|
|
result = SimpleNamespace(
|
|
logits_output=SimpleNamespace(auxiliary_device_output=object()),
|
|
auxiliary_host_output=None,
|
|
copy_done=None,
|
|
copy_to_cpu=Mock(),
|
|
)
|
|
batch = SimpleNamespace(return_logprob=False, return_hidden_states=False)
|
|
|
|
with published_topology():
|
|
Scheduler._copy_auxiliary_output_to_cpu(scheduler, batch, result)
|
|
|
|
assert result.copy_done is event
|
|
result.copy_to_cpu.assert_called_once_with(
|
|
return_logprob=False,
|
|
return_hidden_states=False,
|
|
)
|
|
|
|
|
|
def test_scheduler_preserves_pipeline_parallel_output_for_transport():
|
|
scheduler = object.__new__(Scheduler)
|
|
scheduler.device_module = SimpleNamespace(Event=Mock())
|
|
result = SimpleNamespace(
|
|
logits_output=SimpleNamespace(auxiliary_device_output=object()),
|
|
auxiliary_host_output=None,
|
|
copy_done=None,
|
|
copy_to_cpu=Mock(),
|
|
)
|
|
batch = SimpleNamespace(return_logprob=False, return_hidden_states=False)
|
|
|
|
with published_topology(pp_size=2):
|
|
Scheduler._copy_auxiliary_output_to_cpu(scheduler, batch, result)
|
|
|
|
assert result.copy_done is None
|
|
result.copy_to_cpu.assert_not_called()
|
|
scheduler.device_module.Event.assert_not_called()
|
|
|
|
|
|
def test_pdmux_split_prefill_schedules_auxiliary_output_copy():
|
|
device_output = DeviceOutput(torch.tensor([1.0]))
|
|
result = GenerationBatchResult(
|
|
logits_output=LogitsProcessorOutput(
|
|
next_token_logits=None,
|
|
auxiliary_device_output=device_output,
|
|
),
|
|
next_token_ids=torch.tensor([7]),
|
|
)
|
|
copy_done = CopyDone()
|
|
scheduler = object.__new__(Scheduler)
|
|
scheduler.scheduler_stage_metrics = None
|
|
scheduler.metrics_reporter = Mock()
|
|
scheduler.forward_ct = 0
|
|
scheduler.processed_tokens_counter = 0
|
|
scheduler._sched_idled = False
|
|
scheduler.scripted_scheduler_hook = None
|
|
scheduler.profiler_manager = SimpleNamespace(_profile_batch_predicate=Mock())
|
|
scheduler.forward_sleep_time = None
|
|
scheduler.disaggregation_mode = None
|
|
scheduler.is_generation = True
|
|
scheduler.enable_overlap = False
|
|
scheduler.enable_pdmux = True
|
|
scheduler.tp_worker = SimpleNamespace(
|
|
forward_batch_split_prefill=Mock(return_value=result)
|
|
)
|
|
scheduler.future_map = object()
|
|
scheduler._relay_forward_payload = Mock()
|
|
scheduler.device_module = SimpleNamespace(Event=Mock(return_value=copy_done))
|
|
scheduler.enable_dp_attention = False
|
|
batch = SimpleNamespace(
|
|
forward_mode=SimpleNamespace(
|
|
is_prebuilt=lambda: False,
|
|
is_split_prefill=lambda: True,
|
|
),
|
|
split_index=0,
|
|
reqs=[],
|
|
req_pool_indices=torch.tensor([3]),
|
|
input_ids=torch.tensor([5]),
|
|
extend_num_tokens=1,
|
|
return_logprob=False,
|
|
return_hidden_states=False,
|
|
)
|
|
|
|
with (
|
|
published_topology(),
|
|
patch(
|
|
"sglang.srt.managers.scheduler.resolve_forward_inputs"
|
|
) as resolve_forward_inputs,
|
|
):
|
|
output_result = Scheduler.run_batch(scheduler, batch)
|
|
|
|
resolve_forward_inputs.assert_called_once_with(batch, scheduler.future_map)
|
|
assert output_result is result
|
|
assert result.auxiliary_host_output.values.tolist() == [1.0]
|
|
assert copy_done.record_count == 1
|
|
|
|
|
|
def test_disaggregated_prefill_consumes_auxiliary_output_after_commit():
|
|
host_output = HostOutput(torch.tensor([1.0]))
|
|
copy_done = SimpleNamespace(synchronize=Mock())
|
|
result = GenerationBatchResult(
|
|
logits_output=None,
|
|
next_token_ids=torch.tensor([7]),
|
|
next_draft_input=None,
|
|
copy_done=copy_done,
|
|
auxiliary_host_output=host_output,
|
|
)
|
|
req = SimpleNamespace(
|
|
output_ids=[],
|
|
finished_len=None,
|
|
to_finish=None,
|
|
finished_reason=None,
|
|
inflight_middle_chunks=0,
|
|
pending_bootstrap=False,
|
|
return_logprob=False,
|
|
return_sampling_mask=False,
|
|
grammar=None,
|
|
time_stats=SimpleNamespace(
|
|
set_prefill_finished_time=Mock(),
|
|
set_prefill_transfer_queue_entry_time=Mock(),
|
|
),
|
|
)
|
|
batch = SimpleNamespace(
|
|
reqs=[req],
|
|
spec_info=None,
|
|
prefill_stats=None,
|
|
dp_cooperation_info=None,
|
|
)
|
|
snapshot_auxiliary_output_starts = Mock(
|
|
side_effect=SchedulerBatchResultProcessor.snapshot_auxiliary_output_starts
|
|
)
|
|
processor = SimpleNamespace(
|
|
move_logprobs_to_cpu=Mock(),
|
|
consume_auxiliary_output=Mock(),
|
|
snapshot_auxiliary_output_starts=snapshot_auxiliary_output_starts,
|
|
)
|
|
scheduler = SimpleNamespace(
|
|
batch_result_processor=processor,
|
|
spec_algorithm=SimpleNamespace(is_eagle=lambda: False),
|
|
tree_cache=object(),
|
|
disagg_prefill_inflight_queue=[],
|
|
send_kv_chunk=Mock(),
|
|
metrics_reporter=SimpleNamespace(report_prefill_stats=Mock()),
|
|
maybe_send_health_check_signal=Mock(),
|
|
)
|
|
|
|
with patch("sglang.srt.disaggregation.prefill.maybe_cache_unfinished_req"):
|
|
SchedulerDisaggregationPrefillMixin.process_batch_result_disagg_prefill(
|
|
scheduler,
|
|
batch,
|
|
result,
|
|
)
|
|
|
|
assert req.output_ids == [7]
|
|
snapshot_auxiliary_output_starts.assert_called_once_with(batch, result)
|
|
processor.consume_auxiliary_output.assert_called_once_with(
|
|
batch,
|
|
host_output,
|
|
[0],
|
|
)
|
|
scheduler.maybe_send_health_check_signal.assert_called_once_with()
|
|
|
|
|
|
def test_logprob_only_reuses_preprocessing_without_observer_lifecycle():
|
|
runner = object.__new__(ModelRunner)
|
|
runner._preprocess_logits = Mock()
|
|
runner.sampler = SimpleNamespace(compute_logprobs_only=Mock())
|
|
logits_output = LogitsProcessorOutput(next_token_logits=torch.zeros(1, 4))
|
|
sampling_info = object()
|
|
forward_batch = SimpleNamespace(
|
|
sampling_info=sampling_info,
|
|
top_logprobs_nums=None,
|
|
token_ids_logprobs=[1],
|
|
)
|
|
|
|
ModelRunner.compute_logprobs_only(runner, logits_output, forward_batch)
|
|
|
|
runner._preprocess_logits.assert_called_once_with(logits_output, sampling_info)
|
|
runner.sampler.compute_logprobs_only.assert_called_once()
|
|
|
|
|
|
def test_logprob_only_clears_stale_output_before_early_return():
|
|
runner = object.__new__(ModelRunner)
|
|
runner.sampler = SimpleNamespace(compute_logprobs_only=Mock())
|
|
logits_output = LogitsProcessorOutput(
|
|
next_token_logits=None,
|
|
auxiliary_device_output=DeviceOutput(torch.tensor([99.0])),
|
|
)
|
|
forward_batch = SimpleNamespace(token_ids_logprobs=None)
|
|
|
|
ModelRunner.compute_logprobs_only(runner, logits_output, forward_batch)
|
|
|
|
assert logits_output.auxiliary_device_output is None
|
|
runner.sampler.compute_logprobs_only.assert_not_called()
|
|
|
|
|
|
def test_pipeline_parallel_auxiliary_output_round_trip():
|
|
device_output = DeviceOutput(torch.tensor([1.0, 2.0]))
|
|
result = GenerationBatchResult(
|
|
logits_output=LogitsProcessorOutput(
|
|
next_token_logits=None,
|
|
auxiliary_device_output=device_output,
|
|
),
|
|
next_token_ids=torch.tensor([7]),
|
|
)
|
|
batch = SimpleNamespace(
|
|
return_logprob=False,
|
|
req_pool_indices=torch.tensor([3]),
|
|
input_ids=torch.tensor([5]),
|
|
)
|
|
|
|
tensors = Scheduler._pp_prepare_tensor_dict(
|
|
object.__new__(Scheduler), result, batch
|
|
)
|
|
observer = Observer()
|
|
receiver = object.__new__(Scheduler)
|
|
receiver.pp_group = SimpleNamespace(is_first_rank=True)
|
|
receiver.tp_worker = SimpleNamespace(
|
|
model_runner=SimpleNamespace(sampling_observer=observer)
|
|
)
|
|
receiver.future_map = SimpleNamespace(stash=Mock())
|
|
|
|
output_result = Scheduler._pp_prep_batch_result(
|
|
receiver,
|
|
batch,
|
|
PPBatchMetadata(can_run_cuda_graph=True),
|
|
PPProxyTensors(tensors),
|
|
)
|
|
|
|
assert set(observer.received_tensors) == {"values"}
|
|
assert torch.equal(observer.received_tensors["values"], device_output.values)
|
|
assert output_result.logits_output.auxiliary_device_output is not device_output
|
|
assert torch.equal(output_result.auxiliary_host_output.values, device_output.values)
|
|
assert all("sampling_observer_output" not in key for key in tensors)
|
|
receiver.future_map.stash.assert_called_once()
|
|
|
|
|
|
@pytest.mark.parametrize("dsa_topk_indices", [None, torch.tensor([[2, 5, 7]])])
|
|
def test_pipeline_parallel_dsa_seed_round_trip(dsa_topk_indices):
|
|
draft_input = EagleDraftInput(
|
|
topk_p=torch.tensor([[0.8, 0.2]]),
|
|
topk_index=torch.tensor([[11, 13]]),
|
|
hidden_states=torch.tensor([[1.0, 2.0]]),
|
|
dsa_topk_indices=dsa_topk_indices,
|
|
)
|
|
result = GenerationBatchResult(
|
|
logits_output=None,
|
|
next_token_ids=torch.tensor([17]),
|
|
next_draft_input=draft_input,
|
|
)
|
|
batch = SimpleNamespace(
|
|
return_logprob=False,
|
|
req_pool_indices=torch.tensor([3]),
|
|
input_ids=torch.tensor([5]),
|
|
spec_info=None,
|
|
)
|
|
|
|
tensors = Scheduler._pp_prepare_tensor_dict(
|
|
object.__new__(Scheduler), result, batch
|
|
)
|
|
if dsa_topk_indices is None:
|
|
assert "draft_dsa_topk_indices" not in tensors
|
|
else:
|
|
assert torch.equal(tensors["draft_dsa_topk_indices"], dsa_topk_indices)
|
|
|
|
receiver = object.__new__(Scheduler)
|
|
receiver.pp_group = SimpleNamespace(is_first_rank=False)
|
|
receiver.future_map = SimpleNamespace(stash=Mock())
|
|
Scheduler._pp_prep_batch_result(
|
|
receiver,
|
|
batch,
|
|
PPBatchMetadata(can_run_cuda_graph=True),
|
|
PPProxyTensors(tensors),
|
|
)
|
|
|
|
if dsa_topk_indices is None:
|
|
assert batch.spec_info.dsa_topk_indices is None
|
|
else:
|
|
assert torch.equal(batch.spec_info.dsa_topk_indices, dsa_topk_indices)
|
|
|
|
|
|
def test_pipeline_parallel_auxiliary_output_stays_packed_before_first_rank():
|
|
device_output = DeviceOutput(torch.tensor([1.0]))
|
|
result = GenerationBatchResult(
|
|
logits_output=LogitsProcessorOutput(
|
|
next_token_logits=None,
|
|
auxiliary_device_output=device_output,
|
|
),
|
|
next_token_ids=torch.tensor([7]),
|
|
)
|
|
batch = SimpleNamespace(
|
|
return_logprob=False,
|
|
req_pool_indices=torch.tensor([3]),
|
|
input_ids=torch.tensor([5]),
|
|
)
|
|
tensors = Scheduler._pp_prepare_tensor_dict(
|
|
object.__new__(Scheduler), result, batch
|
|
)
|
|
receiver = object.__new__(Scheduler)
|
|
receiver.pp_group = SimpleNamespace(is_first_rank=False)
|
|
receiver.future_map = SimpleNamespace(stash=Mock())
|
|
|
|
output_result = Scheduler._pp_prep_batch_result(
|
|
receiver,
|
|
batch,
|
|
PPBatchMetadata(can_run_cuda_graph=False),
|
|
PPProxyTensors(tensors),
|
|
)
|
|
|
|
assert output_result.logits_output is None
|
|
assert any("sampling_observer_output" in key for key in tensors)
|
|
|
|
|
|
def test_pipeline_parallel_auxiliary_output_requires_receiver_observer():
|
|
device_output = DeviceOutput(torch.tensor([1.0]))
|
|
result = GenerationBatchResult(
|
|
logits_output=LogitsProcessorOutput(
|
|
next_token_logits=None,
|
|
auxiliary_device_output=device_output,
|
|
),
|
|
next_token_ids=torch.tensor([7]),
|
|
)
|
|
batch = SimpleNamespace(return_logprob=False)
|
|
tensors = Scheduler._pp_prepare_tensor_dict(
|
|
object.__new__(Scheduler), result, batch
|
|
)
|
|
receiver = object.__new__(Scheduler)
|
|
receiver.pp_group = SimpleNamespace(is_first_rank=True)
|
|
receiver.tp_worker = SimpleNamespace(
|
|
model_runner=SimpleNamespace(sampling_observer=None)
|
|
)
|
|
|
|
with pytest.raises(RuntimeError, match="without a sampling observer"):
|
|
Scheduler._pp_prep_batch_result(
|
|
receiver,
|
|
batch,
|
|
PPBatchMetadata(can_run_cuda_graph=False),
|
|
PPProxyTensors(tensors),
|
|
)
|
|
|
|
|
|
def test_pipeline_parallel_auxiliary_output_requires_transport_support():
|
|
result = GenerationBatchResult(
|
|
logits_output=LogitsProcessorOutput(
|
|
next_token_logits=None,
|
|
auxiliary_device_output=HostOnlyDeviceOutput(torch.tensor([1.0])),
|
|
),
|
|
next_token_ids=torch.tensor([7]),
|
|
)
|
|
batch = SimpleNamespace(return_logprob=False)
|
|
|
|
with pytest.raises(RuntimeError, match="does not support pipeline-parallel"):
|
|
Scheduler._pp_prepare_tensor_dict(object.__new__(Scheduler), result, batch)
|
|
|
|
|
|
def test_pipeline_parallel_auxiliary_output_requires_transport_observer():
|
|
result = GenerationBatchResult(
|
|
logits_output=LogitsProcessorOutput(
|
|
next_token_logits=None,
|
|
auxiliary_device_output=DeviceOutput(torch.tensor([1.0])),
|
|
),
|
|
next_token_ids=torch.tensor([7]),
|
|
)
|
|
batch = SimpleNamespace(return_logprob=False)
|
|
tensors = Scheduler._pp_prepare_tensor_dict(
|
|
object.__new__(Scheduler), result, batch
|
|
)
|
|
receiver = object.__new__(Scheduler)
|
|
receiver.pp_group = SimpleNamespace(is_first_rank=True)
|
|
receiver.tp_worker = SimpleNamespace(
|
|
model_runner=SimpleNamespace(sampling_observer=SimpleNamespace())
|
|
)
|
|
|
|
with pytest.raises(RuntimeError, match="does not support pipeline-parallel"):
|
|
Scheduler._pp_prep_batch_result(
|
|
receiver,
|
|
batch,
|
|
PPBatchMetadata(can_run_cuda_graph=False),
|
|
PPProxyTensors(tensors),
|
|
)
|
|
|
|
|
|
def test_auxiliary_output_snapshot_uses_visible_request_lengths():
|
|
batch = SimpleNamespace(
|
|
reqs=[
|
|
SimpleNamespace(output_ids=[10, 11, 12], finished_len=2),
|
|
SimpleNamespace(output_ids=[20], finished_len=None),
|
|
]
|
|
)
|
|
result = SimpleNamespace(auxiliary_host_output=None)
|
|
|
|
assert (
|
|
SchedulerBatchResultProcessor.snapshot_auxiliary_output_starts(batch, result)
|
|
is None
|
|
)
|
|
|
|
result.auxiliary_host_output = object()
|
|
|
|
assert SchedulerBatchResultProcessor.snapshot_auxiliary_output_starts(
|
|
batch, result
|
|
) == [2, 1]
|
|
|
|
|
|
def test_auxiliary_commit_uses_the_scheduler_visible_prefix():
|
|
req = SimpleNamespace(output_ids=[10, 11, 12], finished_len=2)
|
|
|
|
commits = SchedulerBatchResultProcessor._build_auxiliary_commits(
|
|
SimpleNamespace(reqs=[req]),
|
|
output_starts=[1],
|
|
)
|
|
|
|
assert commits[0].output_index == 1
|
|
assert commits[0].token_ids == (11,)
|
|
|
|
|
|
def test_auxiliary_commit_discards_samples_outside_the_visible_output():
|
|
req = SimpleNamespace(output_ids=[10], finished_len=None)
|
|
|
|
commits = SchedulerBatchResultProcessor._build_auxiliary_commits(
|
|
SimpleNamespace(reqs=[req]),
|
|
output_starts=[1],
|
|
)
|
|
|
|
assert commits == [None]
|
|
|
|
|
|
def test_auxiliary_output_consumes_only_newly_visible_tokens():
|
|
req = SimpleNamespace(output_ids=[10, 11, 12], finished_len=2)
|
|
output = Mock()
|
|
batch = SimpleNamespace(reqs=[req])
|
|
|
|
SchedulerBatchResultProcessor.consume_auxiliary_output(
|
|
batch,
|
|
output,
|
|
output_starts=[1],
|
|
)
|
|
|
|
commits = output.consume.call_args.args[1]
|
|
assert commits[0].output_index == 1
|
|
assert commits[0].token_ids == (11,)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main([__file__, "-v"]))
|