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

[Speculative Decoding] Add native UNO serving support (#37667)

Browse Source 
Co-authored-by: drproduck <drproduck@MacBook-Air-2.local>
Co-authored-by: BBuf <1182563586@qq.com>
This commit is contained in:
Yash Akhauri
2026-09-03 20:08:41 +08:00
committed by GitHub
co-authored by drproduck BBuf
parent 354ed6d66b
commit 2bb25dc18b
51 changed files with 6300 additions and 74 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/spec/test_decode_bookkeeping_ownership.py
+6
View File
@@ -45,6 +45,10 @@ _DFLASH_DECODE = (
"speculative/dflash_info_v2.py",
"DFlashDraftInputV2.prepare_for_decode",
)
_UNO_DECODE = (
"speculative/uno_info.py",
"UnoDraftInput.prepare_for_decode",
)
_RESOLVE = (
"managers/scheduler_components/batch_result_processor.py",
"SchedulerBatchResultProcessor._resolve_spec_v2_tokens",
@@ -71,6 +75,8 @@ _OWNER_SITES = {
# one of these two owners for each speculative decode iteration.
(*_DFLASH_DECODE, "decode_batch_idx"): 1,
(*_DFLASH_DECODE, "evict"): 1,
(*_UNO_DECODE, "decode_batch_idx"): 1,
(*_UNO_DECODE, "evict"): 1,
(
"mem_cache/allocation.py",
"alloc_for_spec_decode",
test/registered/unit/spec/test_suffix_attention_merge_dispatch.py
+72
View File
@@ -0,0 +1,72 @@
"""CPU contracts for the fused suffix-attention merge dispatch guard."""
import unittest
from types import SimpleNamespace
import torch
from sglang.kernels.ops.attention.suffix_attention_merge import (
can_use_fused_suffix_attention_merge,
)
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=1, suite="base-a-test-cpu")
class TestSuffixAttentionMergeDispatch(CustomTestCase):
def _inputs(self):
layer = SimpleNamespace(
head_dim=64,
v_head_dim=64,
is_cross_attention=False,
logit_cap=0.0,
)
q = torch.empty((16, 8 * 64), dtype=torch.bfloat16)
key_cache = torch.empty((8, 16, 2, 64), dtype=torch.bfloat16)
value_cache = torch.empty_like(key_cache)
return layer, q, key_cache, value_cache
def _eligible(self, **overrides):
layer, q, key_cache, value_cache = self._inputs()
arguments = dict(
layer=layer,
q=q,
key_cache=key_cache,
value_cache=value_cache,
extra_kwargs={},
)
arguments.update(overrides)
return can_use_fused_suffix_attention_merge(**arguments)
def test_standard_attention_is_eligible(self):
self.assertTrue(self._eligible())
def test_special_attention_features_fall_back(self):
self.assertFalse(self._eligible(extra_kwargs={"sinks": object()}))
layer, _, _, _ = self._inputs()
layer.is_cross_attention = True
self.assertFalse(self._eligible(layer=layer))
layer, _, _, _ = self._inputs()
layer.logit_cap = 20.0
self.assertFalse(self._eligible(layer=layer))
def test_unsupported_tensor_layout_falls_back(self):
layer, _, _, _ = self._inputs()
layer.v_head_dim = 32
self.assertFalse(self._eligible(layer=layer))
_, q, key_cache, value_cache = self._inputs()
self.assertFalse(
self._eligible(
q=q.float(),
key_cache=key_cache.float(),
value_cache=value_cache.float(),
)
)
if __name__ == "__main__":
unittest.main()
test/registered/unit/spec/test_uno_request_validation.py
+93
View File
@@ -0,0 +1,93 @@
"""Request-admission validation for UNO speculative decoding."""
import unittest
from types import SimpleNamespace
from sglang.srt.speculative.uno_validation import validate_uno_request
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=1, suite="base-a-test-cpu")
def _make_request(**overrides):
sampling_params = SimpleNamespace(
min_p=0.0,
json_schema=None,
regex=None,
ebnf=None,
structural_tag=None,
frequency_penalty=0.0,
presence_penalty=0.0,
repetition_penalty=1.0,
min_new_tokens=0,
logit_bias=None,
)
request = SimpleNamespace(
sampling_params=sampling_params,
grammar=None,
return_logprob=False,
return_hidden_states_mode=SimpleNamespace(need_capture=lambda: False),
custom_logit_processor=None,
lora_id=None,
)
for name, value in overrides.items():
target, field = name.split("__", maxsplit=1)
owner = sampling_params if target == "sampling_params" else request
setattr(owner, field, value)
return request
class TestUnoRequestValidation(CustomTestCase):
def test_supported_request_is_accepted(self):
self.assertIsNone(validate_uno_request(_make_request()))
def test_unsupported_request_features_are_rejected(self):
cases = {
"min_p": ({"sampling_params__min_p": 0.1}, "min_p"),
"grammar": ({"sampling_params__regex": "[0-9]+"}, "grammar"),
"logprobs": ({"request__return_logprob": True}, "logprobs"),
"hidden states": (
{
"request__return_hidden_states_mode": SimpleNamespace(
need_capture=lambda: True
)
},
"return_hidden_states",
),
"frequency penalty": (
{"sampling_params__frequency_penalty": 0.1},
"penalties",
),
"presence penalty": (
{"sampling_params__presence_penalty": 0.1},
"penalties",
),
"repetition penalty": (
{"sampling_params__repetition_penalty": 1.1},
"penalties",
),
"minimum new tokens": (
{"sampling_params__min_new_tokens": 1},
"penalties",
),
"logit bias": (
{"sampling_params__logit_bias": {1: 0.5}},
"logit_bias",
),
"custom processor": (
{"request__custom_logit_processor": "processor"},
"custom logit processors",
),
"public LoRA": ({"request__lora_id": "adapter"}, "LoRA"),
}
for name, (overrides, expected) in cases.items():
with self.subTest(name=name):
error = validate_uno_request(_make_request(**overrides))
self.assertIsNotNone(error)
self.assertIn(expected, error)
if __name__ == "__main__":
unittest.main()
test/registered/unit/spec/test_uno_tree_config.py
+73
View File
@@ -0,0 +1,73 @@
"""Startup validation for UNO configuration."""
import unittest
from types import SimpleNamespace
from unittest.mock import patch
from sglang.srt.arg_groups.speculative_hook import _handle_uno
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=1, suite="base-a-test-cpu")
class TestUnoTreeConfig(CustomTestCase):
def test_unsupported_runtime_modes_are_rejected_at_startup(self):
cases = {
"deterministic inference": (
{"enable_deterministic_inference": True},
"enable-deterministic-inference",
),
"strict thinking": (
{"enable_strict_thinking": True},
"enable-strict-thinking",
),
}
for name, (overrides, expected) in cases.items():
with self.subTest(name=name):
values = {
"device": "cuda",
"speculative_draft_model_path": None,
"uno_lora_path": "/tmp/uno-lora",
"enable_deterministic_inference": False,
"enable_strict_thinking": False,
}
values.update(overrides)
server_args = SimpleNamespace(**values)
with (
patch(
"sglang.srt.arg_groups.speculative_hook.resolving_view",
side_effect=lambda args: args,
),
self.assertRaisesRegex(ValueError, expected),
):
_handle_uno(server_args)
def test_parent_list_overflow_is_rejected_at_startup(self):
"""An invalid tree must not survive startup and crash on first decode."""
server_args = SimpleNamespace(
device="cuda",
enable_deterministic_inference=False,
enable_strict_thinking=False,
speculative_draft_model_path=None,
uno_lora_path="/tmp/uno-lora",
speculative_num_draft_tokens=8,
speculative_num_steps=3,
speculative_eagle_topk=2,
)
with (
patch(
"sglang.srt.arg_groups.speculative_hook.resolving_view",
side_effect=lambda args: args,
),
patch("sglang.srt.arg_groups.speculative_hook.declare_resolution"),
self.assertRaisesRegex(ValueError, "parent-list ABI"),
):
_handle_uno(server_args)
if __name__ == "__main__":
unittest.main()
test/registered/unit/spec/test_uno_tree_sparse_sampling.py
+105
View File
@@ -0,0 +1,105 @@
"""CPU contracts for UNO tree compact target sampling."""
import unittest
from types import SimpleNamespace
from unittest.mock import patch
import torch
from sglang.srt.speculative.eagle_utils import (
_can_use_sparse_uno_tree_target_sampling,
)
from sglang.srt.speculative.uno_utils import sample_uno_tree_target_tokens
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=1, suite="base-a-test-cpu")
class TestUnoTreeSparseSampling(CustomTestCase):
def test_sparse_dispatch_guard(self):
sampling_info = SimpleNamespace(
sampling_seed=None,
need_min_p_sampling=False,
)
spec_config = SimpleNamespace(
speculative_use_rejection_sampling=False,
)
with (
patch("sglang.srt.speculative.eagle_utils._is_cuda", True),
patch(
"sglang.srt.speculative.eagle_utils.get_spec",
return_value=spec_config,
),
):
self.assertTrue(
_can_use_sparse_uno_tree_target_sampling(128, sampling_info)
)
self.assertFalse(
_can_use_sparse_uno_tree_target_sampling(None, sampling_info)
)
self.assertFalse(
_can_use_sparse_uno_tree_target_sampling(129, sampling_info)
)
sampling_info.sampling_seed = torch.tensor([1])
self.assertFalse(
_can_use_sparse_uno_tree_target_sampling(128, sampling_info)
)
sampling_info.sampling_seed = None
sampling_info.need_min_p_sampling = True
self.assertFalse(
_can_use_sparse_uno_tree_target_sampling(128, sampling_info)
)
sampling_info.need_min_p_sampling = False
spec_config.speculative_use_rejection_sampling = True
self.assertFalse(
_can_use_sparse_uno_tree_target_sampling(128, sampling_info)
)
def test_targets_are_sampled_from_compact_support(self):
support_ids = torch.tensor(
[
[[10, 11], [20, 21], [30, 31]],
[[40, 41], [50, 51], [60, 61]],
],
dtype=torch.int64,
)
support_probs = torch.full((2, 3, 2), 0.5)
sampled_offsets = torch.tensor(
[[0], [1], [0], [1], [0], [1]],
dtype=torch.long,
)
sampling_info = SimpleNamespace()
next_token_logits = torch.empty((6, 100))
with (
patch(
"sglang.srt.speculative.uno_utils._build_sparse_target_support",
return_value=(support_ids, support_probs),
) as build_support,
patch(
"sglang.srt.speculative.uno_utils.fast_sample",
return_value=(torch.empty((6, 1)), sampled_offsets),
),
):
targets = sample_uno_tree_target_tokens(
next_token_logits=next_token_logits,
sampling_info=sampling_info,
batch_size=2,
verify_width=3,
max_top_k=2,
)
self.assertEqual(targets.tolist(), [[10, 21, 30], [41, 50, 61]])
build_support.assert_called_once_with(
next_token_logits=next_token_logits,
sampling_info=sampling_info,
batch_size=2,
forward_width=3,
max_top_k=2,
)
if __name__ == "__main__":
unittest.main()