4414db594c
Spec-v2 decode was paying unnecessary scheduler-side copy and replay costs: hidden states were copied back to CPU even when not requested, overlap result D2H ran on the forward stream, and draft graph replay issued several small copies separately. Port the focused upstream optimizations without taking the larger spec-v2 relay refactor: gate hidden-state D2H on return_hidden_states, run result copies on copy_stream with pinned async D2H, use stride arange for draft select_index, and group small draft graph replay copies while keeping the large hidden-state DMA copy separate. Constraint: Current branch carries local GLM/NSA/PD/CP changes, so upstream spec-v2 large refactors are not safe to cherry-pick wholesale. Rejected: Cherry-pick upstream spec-v2 relay/dataclass refactor | too broad for this performance fix and conflicts with current branch structure Rejected: Copy hidden states unconditionally | default chat output does not consume them after draft extend Confidence: medium Scope-risk: moderate Directive: Keep result-copy lifetime tied to copy_done; do not move copy_to_cpu back onto forward_stream without measuring decode throughput. Tested: local py_compile for changed production files and new test Tested: local git diff --check Tested: remote g0034 cjy-glm5-new as ubuntu on /sgl-workspace/sglang-tai: pytest -q -p no:cacheprovider test_generation_batch_result_copy.py test_eagle_worker_v2_cp_hidden.py test_spec_utils.py Not-tested: full two-node decode throughput A/B after restart
36 lines
991 B
Python
36 lines
991 B
Python
import torch
|
|
|
|
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
|
|
from sglang.srt.managers.utils import GenerationBatchResult
|
|
|
|
|
|
class _DummyEvent:
|
|
def __init__(self):
|
|
self.recorded = False
|
|
|
|
def record(self):
|
|
self.recorded = True
|
|
|
|
|
|
class _HiddenStateSentinel:
|
|
def to(self, *args, **kwargs):
|
|
raise AssertionError("hidden_states should not be copied when disabled")
|
|
|
|
|
|
def test_copy_to_cpu_skips_hidden_states_when_not_requested():
|
|
hidden_states = _HiddenStateSentinel()
|
|
copy_done = _DummyEvent()
|
|
result = GenerationBatchResult(
|
|
logits_output=LogitsProcessorOutput(
|
|
next_token_logits=None,
|
|
hidden_states=hidden_states,
|
|
),
|
|
next_token_ids=torch.tensor([1], dtype=torch.int64),
|
|
)
|
|
result.copy_done = copy_done
|
|
|
|
result.copy_to_cpu(return_logprob=False, return_hidden_states=False)
|
|
|
|
assert result.logits_output.hidden_states is hidden_states
|
|
assert copy_done.recorded
|