[Sampling] Support sampling masks with overlap scheduling (#36631)
Co-authored-by: ByronHsu <ByronHsu@users.noreply.github.com> Co-authored-by: root <root@slurm-h200-208-179.slurm-compute.tenant-slurm.svc.cluster.local> Co-authored-by: Byron Hsu <byron+per@periodiclabs.ai>
This commit is contained in:
Byron Hsu
co-authored by
ByronHsu
root
Byron Hsu
ByronHsu
root
Byron Hsu
parent
55b45cb45a
commit
fd7743e0e1
@@ -1,3 +1,4 @@
|
||||
import json
|
||||
import math
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
@@ -7,8 +8,14 @@ import requests
|
||||
import torch
|
||||
|
||||
from sglang.srt.layers import sampler as sampler_module
|
||||
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
|
||||
from sglang.srt.layers.sampler import Sampler
|
||||
from sglang.srt.layers.logits_processor import (
|
||||
LogitsProcessorOutput,
|
||||
SamplingMaskStatus,
|
||||
)
|
||||
from sglang.srt.layers.sampler import Sampler, _SamplingMaskCapture
|
||||
from sglang.srt.managers.scheduler_components.batch_result_processor import (
|
||||
SchedulerBatchResultProcessor,
|
||||
)
|
||||
from sglang.srt.sampling.custom_logit_processor import (
|
||||
DisallowedTokensLogitsProcessor,
|
||||
Qwen3ThinkingBudgetLogitProcessor,
|
||||
@@ -24,7 +31,7 @@ from sglang.test.test_utils import (
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
register_cuda_ci(est_time=139, stage="base-b", runner_config="1-gpu-small")
|
||||
register_cuda_ci(est_time=250, stage="base-b", runner_config="2-gpu-large")
|
||||
register_amd_ci(est_time=320, suite="stage-b-test-1-gpu-small-amd")
|
||||
|
||||
_MAX_NEW_TOKENS = 4
|
||||
@@ -36,9 +43,11 @@ _SERVER_ARGS = (
|
||||
"--mem-fraction-static",
|
||||
"0.7",
|
||||
"--enable-custom-logit-processor",
|
||||
"--sampling-mask-max-tokens",
|
||||
"64",
|
||||
)
|
||||
_INVALID_SAMPLING_MASK_ERROR = (
|
||||
"top_p-only sampling is valid but can return huge masks in the tail"
|
||||
"return_sampling_mask requires top_k=1 for greedy sampling"
|
||||
)
|
||||
|
||||
|
||||
@@ -46,6 +55,88 @@ class TestSamplingMaskCapture(CustomTestCase):
|
||||
def setUp(self):
|
||||
self.sampler = Sampler.__new__(Sampler)
|
||||
torch.nn.Module.__init__(self.sampler)
|
||||
self.sampler.sampling_mask_max_tokens = 4096
|
||||
self.sampler.tp_sync_group = None
|
||||
self.sampler.cp_sync_group = None
|
||||
|
||||
def test_default_sampling_does_not_construct_capture_helpers(self):
|
||||
"""Requests without masks must bypass capture-only allocations."""
|
||||
probs = torch.tensor([[0.6, 0.4]])
|
||||
info = SimpleNamespace(sampling_mask_batch_indices=None, sampling_seed=None)
|
||||
with patch.object(
|
||||
sampler_module, "partial", side_effect=AssertionError("capture helper")
|
||||
):
|
||||
_, capture = self.sampler._sample_from_probs(
|
||||
probs=probs,
|
||||
sampling_info=info,
|
||||
positions=torch.tensor([0]),
|
||||
simple_sampling_case=True,
|
||||
)
|
||||
self.assertIsNone(capture)
|
||||
|
||||
def _sample(
|
||||
self, probs, backend, *, top_k=2, top_p=0.45, min_p=0.0, requested_rows=None
|
||||
):
|
||||
batch_size = len(probs)
|
||||
if requested_rows is None:
|
||||
requested_rows = range(batch_size)
|
||||
sampling_info = SimpleNamespace(
|
||||
sampling_seed=None,
|
||||
need_top_k_sampling=True,
|
||||
need_top_p_sampling=top_p < 1.0,
|
||||
need_min_p_sampling=min_p > 0.0,
|
||||
top_ks=torch.full((batch_size,), top_k, dtype=torch.int32, device="cuda"),
|
||||
top_ps=torch.full((batch_size,), top_p, device="cuda"),
|
||||
min_ps=torch.full((batch_size,), min_p, device="cuda"),
|
||||
sampling_mask_batch_indices=torch.tensor(requested_rows, device="cuda"),
|
||||
)
|
||||
with patch.object(
|
||||
sampler_module,
|
||||
"get_exec",
|
||||
return_value=SimpleNamespace(
|
||||
kernel=SimpleNamespace(sampling_backend=backend)
|
||||
),
|
||||
):
|
||||
return self.sampler._sample_from_probs(
|
||||
probs,
|
||||
sampling_info,
|
||||
positions=torch.zeros(batch_size, dtype=torch.int64, device="cuda"),
|
||||
simple_sampling_case=False,
|
||||
)
|
||||
|
||||
def _materialize(self, sampled, capture, requested_rows):
|
||||
output = LogitsProcessorOutput(
|
||||
next_token_logits=None,
|
||||
sampling_mask_output=self.sampler._build_sampling_mask_output(
|
||||
sampled, capture
|
||||
),
|
||||
)
|
||||
output.sampling_mask_output.map_device_tensors(lambda tensor: tensor.cpu())
|
||||
SchedulerBatchResultProcessor.materialize_sampling_mask_output(
|
||||
[
|
||||
SimpleNamespace(return_sampling_mask=i in requested_rows)
|
||||
for i in range(len(sampled))
|
||||
],
|
||||
output,
|
||||
)
|
||||
return output
|
||||
|
||||
def test_min_p_capture_matches_filtered_support_and_logprob(self):
|
||||
backends = ("pytorch",) if is_hip() else ("pytorch", "flashinfer")
|
||||
for backend in backends:
|
||||
with self.subTest(backend=backend):
|
||||
probs = torch.tensor([[0.4, 0.3, 0.2, 0.1]], device="cuda")
|
||||
sampled, capture = self._sample(
|
||||
probs, backend, top_k=3, top_p=1.0, min_p=0.6
|
||||
)
|
||||
output = self.sampler._build_sampling_mask_output(sampled, capture)
|
||||
self.assertEqual(output.statuses.tolist(), [SamplingMaskStatus.OK])
|
||||
self.assertEqual(output.lengths.tolist(), [2])
|
||||
self.assertEqual(set(output.token_ids[0, :2].tolist()), {0, 1})
|
||||
expected = (0.4 if sampled.item() == 0 else 0.3) / 0.7
|
||||
self.assertAlmostEqual(
|
||||
output.selected_logprobs.item(), math.log(expected), places=6
|
||||
)
|
||||
|
||||
def test_hard_exclusion_replay_in_mixed_batch(self):
|
||||
backends = ["pytorch"] if is_hip() else ["pytorch", "flashinfer"]
|
||||
@@ -77,6 +168,7 @@ class TestSamplingMaskCapture(CustomTestCase):
|
||||
)
|
||||
},
|
||||
return_sampling_masks=[True, True],
|
||||
sampling_mask_batch_indices=torch.tensor([0, 1], device="cuda"),
|
||||
)
|
||||
logits = self.sampler._preprocess_logits(logits, info)
|
||||
with patch(
|
||||
@@ -90,12 +182,8 @@ class TestSamplingMaskCapture(CustomTestCase):
|
||||
info,
|
||||
positions=torch.zeros(2, dtype=torch.int64, device="cuda"),
|
||||
simple_sampling_case=False,
|
||||
return_sampling_mask=True,
|
||||
)
|
||||
output = LogitsProcessorOutput(next_token_logits=None)
|
||||
self.sampler._attach_sampling_mask_to_output(
|
||||
output, info, sampled, capture
|
||||
)
|
||||
output = self._materialize(sampled, capture, requested_rows=[0, 1])
|
||||
support = output.next_token_sampling_mask_idx[0]
|
||||
self.assertEqual(set(support), {0, 1, 3})
|
||||
self.assertIn(int(sampled[0]), support)
|
||||
@@ -126,29 +214,7 @@ class TestSamplingMaskCapture(CustomTestCase):
|
||||
expected_ids = expected_support.nonzero(as_tuple=True)[0].tolist()
|
||||
self.assertEqual(expected_ids, [0, 1, 2])
|
||||
|
||||
sampling_info = SimpleNamespace(
|
||||
sampling_seed=None,
|
||||
need_top_k_sampling=True,
|
||||
need_top_p_sampling=True,
|
||||
need_min_p_sampling=False,
|
||||
top_ks=torch.full((batch_size,), top_k, dtype=torch.int32, device="cuda"),
|
||||
top_ps=torch.full((batch_size,), top_p, device="cuda"),
|
||||
min_ps=torch.zeros(batch_size, device="cuda"),
|
||||
return_sampling_masks=[True] * batch_size,
|
||||
)
|
||||
with patch(
|
||||
"sglang.srt.layers.sampler.get_exec",
|
||||
return_value=SimpleNamespace(
|
||||
kernel=SimpleNamespace(sampling_backend="flashinfer")
|
||||
),
|
||||
):
|
||||
sampled, capture = self.sampler._sample_from_probs(
|
||||
probs,
|
||||
sampling_info,
|
||||
positions=torch.zeros(batch_size, dtype=torch.int64, device="cuda"),
|
||||
simple_sampling_case=False,
|
||||
return_sampling_mask=True,
|
||||
)
|
||||
sampled, capture = self._sample(probs, "flashinfer", top_k=top_k, top_p=top_p)
|
||||
|
||||
self.assertIsNotNone(capture)
|
||||
self.assertEqual(capture.batch_rows.cpu().tolist(), list(range(batch_size)))
|
||||
@@ -164,46 +230,24 @@ class TestSamplingMaskCapture(CustomTestCase):
|
||||
@unittest.skipIf(is_hip(), "FlashInfer is not available on ROCm")
|
||||
def test_flashinfer_capture_only_materializes_requested_rows(self):
|
||||
batch_size = 4
|
||||
top_k = 2
|
||||
top_p = 0.45
|
||||
requested_rows = [1, 3]
|
||||
probs = torch.tensor([[0.4, 0.2, 0.2, 0.1, 0.1]], device="cuda").repeat(
|
||||
batch_size, 1
|
||||
)
|
||||
sampling_info = SimpleNamespace(
|
||||
sampling_seed=None,
|
||||
need_top_k_sampling=True,
|
||||
need_top_p_sampling=True,
|
||||
need_min_p_sampling=False,
|
||||
top_ks=torch.full((batch_size,), top_k, dtype=torch.int32, device="cuda"),
|
||||
top_ps=torch.full((batch_size,), top_p, device="cuda"),
|
||||
min_ps=torch.zeros(batch_size, device="cuda"),
|
||||
return_sampling_masks=[False, True, False, True],
|
||||
)
|
||||
top_k_renorm = sampler_module.top_k_renorm_prob
|
||||
top_p_renorm = sampler_module.top_p_renorm_prob
|
||||
with (
|
||||
patch(
|
||||
"sglang.srt.layers.sampler.get_exec",
|
||||
return_value=SimpleNamespace(
|
||||
kernel=SimpleNamespace(sampling_backend="flashinfer")
|
||||
),
|
||||
),
|
||||
patch(
|
||||
"sglang.srt.layers.sampler.top_k_renorm_prob",
|
||||
wraps=top_k_renorm,
|
||||
patch.object(
|
||||
sampler_module,
|
||||
"top_k_renorm_prob",
|
||||
wraps=sampler_module.top_k_renorm_prob,
|
||||
) as top_k_mock,
|
||||
patch(
|
||||
"sglang.srt.layers.sampler.top_p_renorm_prob",
|
||||
wraps=top_p_renorm,
|
||||
patch.object(
|
||||
sampler_module,
|
||||
"top_p_renorm_prob",
|
||||
wraps=sampler_module.top_p_renorm_prob,
|
||||
) as top_p_mock,
|
||||
):
|
||||
sampled, capture = self.sampler._sample_from_probs(
|
||||
probs,
|
||||
sampling_info,
|
||||
positions=torch.zeros(batch_size, dtype=torch.int64, device="cuda"),
|
||||
simple_sampling_case=False,
|
||||
return_sampling_mask=True,
|
||||
sampled, capture = self._sample(
|
||||
probs, "flashinfer", requested_rows=requested_rows
|
||||
)
|
||||
|
||||
self.assertIsNotNone(capture)
|
||||
@@ -212,10 +256,7 @@ class TestSamplingMaskCapture(CustomTestCase):
|
||||
self.assertEqual(tuple(top_k_mock.call_args.args[0].shape), (2, 5))
|
||||
self.assertEqual(tuple(top_p_mock.call_args.args[0].shape), (2, 5))
|
||||
|
||||
output = LogitsProcessorOutput(next_token_logits=None)
|
||||
self.sampler._attach_sampling_mask_to_output(
|
||||
output, sampling_info, sampled, capture
|
||||
)
|
||||
output = self._materialize(sampled, capture, requested_rows)
|
||||
self.assertIsNone(output.next_token_sampling_mask_idx[0])
|
||||
self.assertEqual(set(output.next_token_sampling_mask_idx[1]), {0, 1, 2})
|
||||
self.assertIsNone(output.next_token_sampling_mask_idx[2])
|
||||
@@ -231,39 +272,14 @@ class TestSamplingMaskCapture(CustomTestCase):
|
||||
probs = torch.tensor([[0.4, 0.2, 0.2, 0.1, 0.1]], device="cuda").repeat(
|
||||
batch_size, 1
|
||||
)
|
||||
sampling_info = SimpleNamespace(
|
||||
sampling_seed=None,
|
||||
need_top_k_sampling=True,
|
||||
need_top_p_sampling=True,
|
||||
need_min_p_sampling=False,
|
||||
top_ks=torch.full((batch_size,), 2, dtype=torch.int32, device="cuda"),
|
||||
top_ps=torch.full((batch_size,), 0.45, device="cuda"),
|
||||
min_ps=torch.zeros(batch_size, device="cuda"),
|
||||
return_sampling_masks=[False, True, False, True],
|
||||
)
|
||||
with patch(
|
||||
"sglang.srt.layers.sampler.get_exec",
|
||||
return_value=SimpleNamespace(
|
||||
kernel=SimpleNamespace(sampling_backend="pytorch")
|
||||
),
|
||||
):
|
||||
sampled, capture = self.sampler._sample_from_probs(
|
||||
probs,
|
||||
sampling_info,
|
||||
positions=torch.zeros(batch_size, dtype=torch.int64, device="cuda"),
|
||||
simple_sampling_case=False,
|
||||
return_sampling_mask=True,
|
||||
)
|
||||
sampled, capture = self._sample(probs, "pytorch", requested_rows=requested_rows)
|
||||
|
||||
self.assertIsNotNone(capture)
|
||||
self.assertEqual(capture.batch_rows.cpu().tolist(), requested_rows)
|
||||
self.assertEqual(tuple(capture.weights.shape), (len(requested_rows), 5))
|
||||
self.assertEqual(tuple(capture.token_ids.shape), (len(requested_rows), 5))
|
||||
|
||||
output = LogitsProcessorOutput(next_token_logits=None)
|
||||
self.sampler._attach_sampling_mask_to_output(
|
||||
output, sampling_info, sampled, capture
|
||||
)
|
||||
output = self._materialize(sampled, capture, requested_rows)
|
||||
for batch_row in requested_rows:
|
||||
self.assertIn(
|
||||
int(sampled[batch_row]),
|
||||
@@ -297,18 +313,38 @@ class SamplingMaskTestMixin:
|
||||
return_logprob=False,
|
||||
top_logprobs_num=0,
|
||||
custom_logit_processor=None,
|
||||
stream=False,
|
||||
):
|
||||
payload = {
|
||||
"text": "The capital of France is",
|
||||
"sampling_params": sampling_params,
|
||||
"sampling_params": {
|
||||
"temperature": 1.0,
|
||||
"max_new_tokens": _MAX_NEW_TOKENS,
|
||||
"ignore_eos": True,
|
||||
**sampling_params,
|
||||
},
|
||||
"return_sampling_mask": return_sampling_mask,
|
||||
"stream": stream,
|
||||
}
|
||||
if custom_logit_processor is not None:
|
||||
payload["custom_logit_processor"] = custom_logit_processor
|
||||
if return_logprob:
|
||||
payload["return_logprob"] = True
|
||||
payload["top_logprobs_num"] = top_logprobs_num
|
||||
return requests.post(self.base_url + "/generate", json=payload, timeout=60)
|
||||
return requests.post(
|
||||
self.base_url + "/generate", json=payload, stream=stream, timeout=60
|
||||
)
|
||||
|
||||
def _assert_sampling_masks(self, output_ids, meta_info):
|
||||
masks = meta_info["output_token_sampling_mask"]
|
||||
self.assertEqual(len(masks), len(output_ids))
|
||||
self.assertEqual(
|
||||
len(meta_info["output_token_sampling_logprobs"]), len(output_ids)
|
||||
)
|
||||
for token_id, mask in zip(output_ids, masks):
|
||||
self.assertIn(token_id, mask)
|
||||
self.assertEqual(len(mask), len(set(mask)))
|
||||
return masks
|
||||
|
||||
def _generate_sampling_masks(self, sampling_params):
|
||||
response = self._post_generate(sampling_params)
|
||||
@@ -317,23 +353,13 @@ class SamplingMaskTestMixin:
|
||||
output = response.json()
|
||||
meta_info = output["meta_info"]
|
||||
output_ids = output["output_ids"]
|
||||
sampling_masks = meta_info["output_token_sampling_mask"]
|
||||
|
||||
self.assertEqual(len(output_ids), _MAX_NEW_TOKENS)
|
||||
self.assertEqual(meta_info["completion_tokens"], len(output_ids))
|
||||
self.assertEqual(
|
||||
meta_info["output_token_sampling_mask_length"], len(output_ids)
|
||||
)
|
||||
self.assertEqual(len(sampling_masks), len(output_ids))
|
||||
for output_id, sampling_mask in zip(output_ids, sampling_masks):
|
||||
self.assertIn(output_id, sampling_mask)
|
||||
self.assertEqual(len(sampling_mask), len(set(sampling_mask)))
|
||||
return sampling_masks
|
||||
|
||||
def _assert_rejects_unbounded_sampling_mask(self, sampling_params):
|
||||
response = self._post_generate(sampling_params)
|
||||
self.assertEqual(response.status_code, 400, response.text)
|
||||
self.assertIn(_INVALID_SAMPLING_MASK_ERROR, response.text)
|
||||
return self._assert_sampling_masks(output_ids, meta_info)
|
||||
|
||||
|
||||
class TestSamplingMask(SamplingMaskTestMixin, CustomTestCase):
|
||||
@@ -393,40 +419,19 @@ class TestSamplingMask(SamplingMaskTestMixin, CustomTestCase):
|
||||
self.assertEqual(recovery.status_code, 200, recovery.text)
|
||||
|
||||
def test_generate_returns_sampling_mask(self):
|
||||
top_p_sampling_masks = self._generate_sampling_masks(
|
||||
{
|
||||
"temperature": 1.0,
|
||||
"top_k": _TOP_K,
|
||||
"top_p": _TOP_P,
|
||||
"max_new_tokens": _MAX_NEW_TOKENS,
|
||||
"ignore_eos": True,
|
||||
}
|
||||
)
|
||||
for sampling_mask in top_p_sampling_masks:
|
||||
self.assertGreater(len(sampling_mask), 0)
|
||||
for params, min_size in (
|
||||
({"top_p": _TOP_P}, 1),
|
||||
({}, _TOP_K),
|
||||
({"top_p": 1.0}, _TOP_K),
|
||||
):
|
||||
with self.subTest(sampling_params=params):
|
||||
masks = self._generate_sampling_masks({"top_k": _TOP_K, **params})
|
||||
for mask in masks:
|
||||
self.assertGreaterEqual(len(mask), min_size)
|
||||
|
||||
top_k_sampling_masks = self._generate_sampling_masks(
|
||||
{
|
||||
"temperature": 1.0,
|
||||
"top_k": _TOP_K,
|
||||
"max_new_tokens": _MAX_NEW_TOKENS,
|
||||
"ignore_eos": True,
|
||||
}
|
||||
)
|
||||
for sampling_mask in top_k_sampling_masks:
|
||||
self.assertGreaterEqual(len(sampling_mask), _TOP_K)
|
||||
|
||||
top_k_top_p_one_sampling_masks = self._generate_sampling_masks(
|
||||
{
|
||||
"temperature": 1.0,
|
||||
"top_k": _TOP_K,
|
||||
"top_p": 1.0,
|
||||
"max_new_tokens": _MAX_NEW_TOKENS,
|
||||
"ignore_eos": True,
|
||||
}
|
||||
)
|
||||
for sampling_mask in top_k_top_p_one_sampling_masks:
|
||||
self.assertGreaterEqual(len(sampling_mask), _TOP_K)
|
||||
def test_generate_returns_greedy_singleton_mask(self):
|
||||
masks = self._generate_sampling_masks({"temperature": 0.0})
|
||||
self.assertTrue(all(len(mask) == 1 for mask in masks))
|
||||
|
||||
def test_sampling_mask_matches_topk_logprobs(self):
|
||||
"""Check the returned mask and its renormalized logprobs.
|
||||
@@ -444,13 +449,7 @@ class TestSamplingMask(SamplingMaskTestMixin, CustomTestCase):
|
||||
"""
|
||||
top_k, top_p = _TOP_K, _TOP_P
|
||||
response = self._post_generate(
|
||||
{
|
||||
"temperature": 1.0,
|
||||
"top_k": top_k,
|
||||
"top_p": top_p,
|
||||
"max_new_tokens": _MAX_NEW_TOKENS,
|
||||
"ignore_eos": True,
|
||||
},
|
||||
{"top_k": top_k, "top_p": top_p},
|
||||
return_logprob=True,
|
||||
top_logprobs_num=_TOP_LOGPROBS_NUM,
|
||||
)
|
||||
@@ -459,12 +458,10 @@ class TestSamplingMask(SamplingMaskTestMixin, CustomTestCase):
|
||||
output = response.json()
|
||||
meta_info = output["meta_info"]
|
||||
output_ids = output["output_ids"]
|
||||
sampling_masks = meta_info["output_token_sampling_mask"]
|
||||
sampling_masks = self._assert_sampling_masks(output_ids, meta_info)
|
||||
sampling_logprobs = meta_info["output_token_sampling_logprobs"]
|
||||
top_logprobs = meta_info["output_top_logprobs"] # [logprob, id, text] per token
|
||||
|
||||
self.assertEqual(len(sampling_masks), len(output_ids))
|
||||
self.assertEqual(len(sampling_logprobs), len(output_ids))
|
||||
self.assertEqual(len(top_logprobs), len(output_ids))
|
||||
|
||||
for output_id, mask, mask_logprob, step_top_logprobs in zip(
|
||||
@@ -476,7 +473,6 @@ class TestSamplingMask(SamplingMaskTestMixin, CustomTestCase):
|
||||
|
||||
mask_set = set(mask)
|
||||
|
||||
self.assertIn(output_id, mask_set)
|
||||
self.assertTrue(mask_set.issubset(probs))
|
||||
top_k_cutoff = sorted(probs.values(), reverse=True)[top_k - 1]
|
||||
for token_id in mask_set:
|
||||
@@ -508,33 +504,115 @@ class TestSamplingMask(SamplingMaskTestMixin, CustomTestCase):
|
||||
|
||||
choice = response.json()["choices"][0]
|
||||
output_ids = choice["response_token_ids"]
|
||||
meta_info = choice["meta_info"]
|
||||
sampling_masks = meta_info["output_token_sampling_mask"]
|
||||
sampling_logprobs = meta_info["output_token_sampling_logprobs"]
|
||||
self.assertEqual(len(output_ids), _MAX_NEW_TOKENS)
|
||||
self._assert_sampling_masks(output_ids, choice["meta_info"])
|
||||
|
||||
def test_generate_streams_aligned_sampling_masks(self):
|
||||
response = self._post_generate({"top_k": _TOP_K, "top_p": _TOP_P}, stream=True)
|
||||
self.assertEqual(response.status_code, 200, response.text)
|
||||
|
||||
output_ids = []
|
||||
for line in response.iter_lines():
|
||||
if not line.startswith(b"data: ") or line[6:] == b"[DONE]":
|
||||
continue
|
||||
chunk = json.loads(line[6:])
|
||||
output_ids = chunk["output_ids"]
|
||||
self._assert_sampling_masks(output_ids, chunk["meta_info"])
|
||||
|
||||
self.assertEqual(len(output_ids), _MAX_NEW_TOKENS)
|
||||
self.assertEqual(len(sampling_masks), len(output_ids))
|
||||
self.assertEqual(len(sampling_logprobs), len(output_ids))
|
||||
for output_id, sampling_mask in zip(output_ids, sampling_masks):
|
||||
self.assertIn(output_id, sampling_mask)
|
||||
|
||||
def test_generate_rejects_unbounded_sampling_mask(self):
|
||||
self._assert_rejects_unbounded_sampling_mask(
|
||||
{
|
||||
"temperature": 1.0,
|
||||
"top_p": _TOP_P,
|
||||
"max_new_tokens": _MAX_NEW_TOKENS,
|
||||
"ignore_eos": True,
|
||||
}
|
||||
for params in ({"top_p": _TOP_P}, {"top_k": 65}, {"top_p": 1.0}):
|
||||
with self.subTest(sampling_params=params):
|
||||
response = self._post_generate(params)
|
||||
self.assertEqual(response.status_code, 400, response.text)
|
||||
self.assertIn(_INVALID_SAMPLING_MASK_ERROR, response.text)
|
||||
|
||||
|
||||
class TestSamplingMaskPacking(CustomTestCase):
|
||||
def setUp(self):
|
||||
self.sampler = Sampler.__new__(Sampler)
|
||||
self.sampler.sampling_mask_max_tokens = 3
|
||||
self.sampler.tp_sync_group = None
|
||||
self.sampler.cp_sync_group = None
|
||||
|
||||
def test_selected_token_must_have_positive_captured_weight(self):
|
||||
for token_ids in (None, torch.tensor([[2, 1, 0]], dtype=torch.int32)):
|
||||
with self.subTest(sorted_capture=token_ids is not None):
|
||||
capture = _SamplingMaskCapture(
|
||||
batch_rows=torch.tensor([0]),
|
||||
weights=torch.tensor([[0.7, 0.3, 0.0]]),
|
||||
token_ids=token_ids,
|
||||
selected_weight=None,
|
||||
)
|
||||
selected = torch.tensor([2 if token_ids is None else 0])
|
||||
output = self.sampler._build_sampling_mask_output(selected, capture)
|
||||
self.assertEqual(output.statuses.tolist(), [SamplingMaskStatus.INVALID])
|
||||
|
||||
def test_synced_token_logprob_is_recomputed_from_capture(self):
|
||||
capture = _SamplingMaskCapture(
|
||||
batch_rows=torch.tensor([0]),
|
||||
weights=torch.tensor([[0.6, 0.2, 0.0]]),
|
||||
token_ids=torch.tensor([[2, 1, 0]], dtype=torch.int32),
|
||||
selected_weight=None,
|
||||
)
|
||||
self._assert_rejects_unbounded_sampling_mask(
|
||||
{
|
||||
"temperature": 1.0,
|
||||
"top_p": 1.0,
|
||||
"max_new_tokens": _MAX_NEW_TOKENS,
|
||||
"ignore_eos": True,
|
||||
}
|
||||
output = self.sampler._build_sampling_mask_output(torch.tensor([1]), capture)
|
||||
self.assertEqual(output.statuses.tolist(), [SamplingMaskStatus.OK])
|
||||
self.assertAlmostEqual(output.selected_logprobs.item(), math.log(0.25))
|
||||
|
||||
def test_greedy_device_output_survives_async_copy(self):
|
||||
from sglang.srt.managers.utils import GenerationBatchResult
|
||||
|
||||
tokens = torch.tensor([3, 4, 5], device="cuda")
|
||||
output = LogitsProcessorOutput(
|
||||
next_token_logits=None,
|
||||
sampling_mask_output=self.sampler._build_greedy_sampling_mask_output(
|
||||
torch.tensor([0, 2], device="cuda"), tokens
|
||||
),
|
||||
)
|
||||
result = GenerationBatchResult(
|
||||
logits_output=output, next_token_ids=tokens, copy_done=torch.cuda.Event()
|
||||
)
|
||||
result.copy_to_cpu(return_logprob=False)
|
||||
result.copy_done.synchronize()
|
||||
self.assertEqual(output.sampling_mask_output.token_ids.device.type, "cpu")
|
||||
SchedulerBatchResultProcessor.materialize_sampling_mask_output(
|
||||
[
|
||||
SimpleNamespace(return_sampling_mask=flag)
|
||||
for flag in (True, False, True)
|
||||
],
|
||||
output,
|
||||
)
|
||||
self.assertEqual(output.next_token_sampling_mask_idx, [[3], None, [5]])
|
||||
self.assertEqual(output.next_token_sampling_logprobs, [0.0, None, 0.0])
|
||||
|
||||
def test_overflow_never_materializes_a_partial_mask(self):
|
||||
# Simulate a top-k cutoff tie: a nominal top_k below the cap can still
|
||||
# produce more positive weights than the fixed transport can hold.
|
||||
capture = _SamplingMaskCapture(
|
||||
batch_rows=torch.tensor([0]),
|
||||
weights=torch.tensor([[0.2, 0.2, 0.2, 0.2, 0.2]]),
|
||||
token_ids=None,
|
||||
selected_weight=torch.tensor([0.2]),
|
||||
)
|
||||
|
||||
sampling_output = self.sampler._build_sampling_mask_output(
|
||||
torch.tensor([0]), capture
|
||||
)
|
||||
|
||||
output = LogitsProcessorOutput(
|
||||
next_token_logits=None,
|
||||
sampling_mask_output=sampling_output,
|
||||
)
|
||||
SchedulerBatchResultProcessor.materialize_sampling_mask_output(
|
||||
[SimpleNamespace(return_sampling_mask=True)], output
|
||||
)
|
||||
self.assertEqual(
|
||||
output.next_token_sampling_mask_status,
|
||||
[SamplingMaskStatus.OVERFLOW],
|
||||
)
|
||||
self.assertEqual(output.next_token_sampling_mask_idx, [None])
|
||||
self.assertEqual(output.next_token_sampling_logprobs, [None])
|
||||
|
||||
|
||||
class TestSamplingMaskDeterministic(SamplingMaskTestMixin, CustomTestCase):
|
||||
@@ -548,32 +626,20 @@ class TestSamplingMaskDeterministic(SamplingMaskTestMixin, CustomTestCase):
|
||||
|
||||
def test_return_sampling_mask_preserves_deterministic_sampling(self):
|
||||
sampling_params = {
|
||||
"temperature": 1.0,
|
||||
"top_k": _TOP_K,
|
||||
"top_p": 1.0,
|
||||
"sampling_seed": _SAMPLING_SEED,
|
||||
"max_new_tokens": _MAX_NEW_TOKENS,
|
||||
"ignore_eos": True,
|
||||
}
|
||||
|
||||
with_mask_response = self._post_generate(
|
||||
sampling_params, return_sampling_mask=True
|
||||
)
|
||||
self.assertEqual(with_mask_response.status_code, 200, with_mask_response.text)
|
||||
|
||||
without_mask_response = self._post_generate(
|
||||
sampling_params, return_sampling_mask=False
|
||||
)
|
||||
self.assertEqual(
|
||||
without_mask_response.status_code, 200, without_mask_response.text
|
||||
)
|
||||
|
||||
with_mask_output = with_mask_response.json()
|
||||
without_mask_output = without_mask_response.json()
|
||||
self.assertEqual(
|
||||
with_mask_output["output_ids"], without_mask_output["output_ids"]
|
||||
)
|
||||
self.assertEqual(with_mask_output["text"], without_mask_output["text"])
|
||||
outputs = []
|
||||
for return_mask in (False, True):
|
||||
response = self._post_generate(
|
||||
sampling_params, return_sampling_mask=return_mask
|
||||
)
|
||||
self.assertEqual(response.status_code, 200, response.text)
|
||||
output = response.json()
|
||||
outputs.append((output["output_ids"], output["text"]))
|
||||
self.assertEqual(outputs[0], outputs[1])
|
||||
|
||||
|
||||
class TestSamplingMaskPytorch(TestSamplingMask):
|
||||
@@ -584,5 +650,88 @@ class TestSamplingMaskPytorch(TestSamplingMask):
|
||||
cls._launch_server(("--sampling-backend", "pytorch"))
|
||||
|
||||
|
||||
@unittest.skipIf(is_hip(), "The AMD sampling-mask CI suite provides only one GPU.")
|
||||
class TestDistributedSamplingMask(CustomTestCase):
|
||||
def _check_parallel_config(self, *, tp_size, pp_size):
|
||||
process = None
|
||||
try:
|
||||
process = popen_launch_server(
|
||||
"Qwen/Qwen2.5-0.5B-Instruct",
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=[
|
||||
"--tp-size",
|
||||
str(tp_size),
|
||||
"--pp-size",
|
||||
str(pp_size),
|
||||
"--sampling-mask-max-tokens",
|
||||
"64",
|
||||
"--mem-fraction-static",
|
||||
"0.5",
|
||||
"--max-running-requests",
|
||||
"8",
|
||||
"--cuda-graph-max-bs-decode",
|
||||
"8",
|
||||
],
|
||||
)
|
||||
for return_logprob in (False, True):
|
||||
with self.subTest(return_logprob=return_logprob):
|
||||
output = self._generate(
|
||||
return_sampling_mask=True, return_logprob=return_logprob
|
||||
)
|
||||
token_ids = output["output_ids"]
|
||||
meta = output["meta_info"]
|
||||
masks = meta["output_token_sampling_mask"]
|
||||
logprobs = meta["output_token_sampling_logprobs"]
|
||||
self.assertEqual(len(token_ids), 4)
|
||||
self.assertEqual(meta["output_token_sampling_mask_length"], 4)
|
||||
self.assertEqual(len(masks), 4)
|
||||
self.assertEqual(len(logprobs), 4)
|
||||
for token_id, mask, logprob in zip(token_ids, masks, logprobs):
|
||||
self.assertIn(token_id, mask)
|
||||
self.assertEqual(len(mask), len(set(mask)))
|
||||
self.assertLessEqual(len(mask), 64)
|
||||
self.assertTrue(math.isfinite(logprob))
|
||||
self.assertLessEqual(logprob, 0.0)
|
||||
if return_logprob:
|
||||
self.assertEqual(len(meta["output_token_logprobs"]), 4)
|
||||
|
||||
ordinary = self._generate(return_sampling_mask=False, return_logprob=False)
|
||||
self.assertEqual(len(ordinary["output_ids"]), 4)
|
||||
self.assertNotIn("output_token_sampling_mask", ordinary["meta_info"])
|
||||
finally:
|
||||
if process is not None:
|
||||
kill_process_tree(process.pid)
|
||||
process.wait(timeout=30)
|
||||
|
||||
def _generate(self, *, return_sampling_mask, return_logprob):
|
||||
response = requests.post(
|
||||
DEFAULT_URL_FOR_TEST + "/generate",
|
||||
json={
|
||||
"text": "The capital of France is",
|
||||
"sampling_params": {
|
||||
"temperature": 0.8,
|
||||
"top_k": 8,
|
||||
"top_p": 0.9,
|
||||
"max_new_tokens": 4,
|
||||
"ignore_eos": True,
|
||||
},
|
||||
"return_sampling_mask": return_sampling_mask,
|
||||
"return_logprob": return_logprob,
|
||||
},
|
||||
timeout=120,
|
||||
)
|
||||
self.assertEqual(response.status_code, 200, response.text)
|
||||
return response.json()
|
||||
|
||||
def test_tp2_sampling_mask(self):
|
||||
"""Exercise status synchronization across two tensor-parallel ranks."""
|
||||
self._check_parallel_config(tp_size=2, pp_size=1)
|
||||
|
||||
def test_pp2_sampling_mask(self):
|
||||
"""Exercise mask transport between two live pipeline stages."""
|
||||
self._check_parallel_config(tp_size=1, pp_size=2)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -46,6 +46,7 @@ from sglang.srt.speculative.eagle_disaggregation import (
|
||||
build_eagle_disagg_draft_input,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cpu_ci(est_time=11, suite="base-a-test-cpu")
|
||||
|
||||
@@ -335,9 +336,14 @@ class TestMooncakePPStaging(unittest.TestCase):
|
||||
)
|
||||
|
||||
|
||||
class TestEagleDsaSeedTransfer(unittest.TestCase):
|
||||
class TestEagleDsaSeedTransfer(CustomTestCase):
|
||||
@staticmethod
|
||||
def _make_req(seed, metadata_buffer_index=0):
|
||||
def _make_req(
|
||||
seed,
|
||||
metadata_buffer_index=0,
|
||||
sampling_mask=None,
|
||||
sampling_logprob=None,
|
||||
):
|
||||
return SimpleNamespace(
|
||||
metadata_buffer_index=metadata_buffer_index,
|
||||
output_ids=[101],
|
||||
@@ -347,7 +353,13 @@ class TestEagleDsaSeedTransfer(unittest.TestCase):
|
||||
cached_tokens_storage=0,
|
||||
multimodal_inputs=None,
|
||||
return_logprob=False,
|
||||
return_sampling_mask=False,
|
||||
return_sampling_mask=sampling_mask is not None,
|
||||
output_token_sampling_mask=(
|
||||
None if sampling_mask is None else [sampling_mask]
|
||||
),
|
||||
output_token_sampling_logprobs=(
|
||||
None if sampling_logprob is None else [sampling_logprob]
|
||||
),
|
||||
hidden_states_tensor=torch.tensor([1.0, 2.0]),
|
||||
output_topk_p=torch.tensor([1.0]),
|
||||
output_topk_index=torch.tensor([7]),
|
||||
@@ -360,6 +372,7 @@ class TestEagleDsaSeedTransfer(unittest.TestCase):
|
||||
size=2,
|
||||
hidden_size=2,
|
||||
hidden_states_dtype=torch.float32,
|
||||
max_sampling_mask_tokens=16,
|
||||
output_dsa_topk_indices_dim=3,
|
||||
)
|
||||
seed = torch.tensor([4, 5, 6], dtype=torch.int32)
|
||||
@@ -378,6 +391,46 @@ class TestEagleDsaSeedTransfer(unittest.TestCase):
|
||||
self.assertEqual(data_lens[-2], buffers.output_dsa_topk_indices.nbytes)
|
||||
self.assertEqual(item_lens[-2], buffers.output_dsa_topk_indices[0].nbytes)
|
||||
|
||||
def test_sampling_mask_metadata_is_opt_in(self):
|
||||
"""Disabled masks stay off the wire; enabled masks round-trip at capacity."""
|
||||
schemas = []
|
||||
for enabled in (False, True):
|
||||
with (
|
||||
self.subTest(enabled=enabled),
|
||||
envs.SGLANG_ENABLE_DISAGG_SAMPLING_MASK.override(enabled),
|
||||
):
|
||||
buffers = MetadataBuffers(
|
||||
size=1,
|
||||
hidden_size=2,
|
||||
hidden_states_dtype=torch.float32,
|
||||
max_sampling_mask_tokens=3,
|
||||
)
|
||||
buffers.set_buf(
|
||||
self._make_req(
|
||||
None,
|
||||
sampling_mask=[7, 8, 9] if enabled else None,
|
||||
sampling_logprob=-1.25 if enabled else None,
|
||||
)
|
||||
)
|
||||
schemas.append(buffers.get_buf_infos())
|
||||
if enabled:
|
||||
self.assertEqual(
|
||||
buffers.output_token_sampling_mask_idx.shape, (1, 3)
|
||||
)
|
||||
length, mask, logprob = buffers.get_buf(0)[6:9]
|
||||
self.assertEqual(length[0].item(), 3)
|
||||
self.assertEqual(mask.tolist(), [7, 8, 9])
|
||||
self.assertAlmostEqual(logprob[0].item(), -1.25)
|
||||
else:
|
||||
self.assertIsNone(buffers.output_token_sampling_mask_len)
|
||||
self.assertIsNone(buffers.output_token_sampling_mask_idx)
|
||||
self.assertIsNone(buffers.output_token_sampling_logprobs)
|
||||
self.assertEqual(buffers.get_buf(0)[6:9], (None, None, None))
|
||||
disabled_ptrs, _, disabled_sizes = schemas[0]
|
||||
enabled_ptrs, _, enabled_sizes = schemas[1]
|
||||
self.assertEqual(len(enabled_ptrs) - len(disabled_ptrs), 3)
|
||||
self.assertEqual(sum(enabled_sizes) - sum(disabled_sizes), 3 * 4 + 128)
|
||||
|
||||
def test_decode_input_requires_valid_seed_for_every_request(self):
|
||||
seeds = (
|
||||
torch.tensor([1, 2, 3], dtype=torch.int32),
|
||||
|
||||
@@ -5,8 +5,13 @@ import pytest
|
||||
import torch
|
||||
|
||||
from sglang.srt.disaggregation.prefill import SchedulerDisaggregationPrefillMixin
|
||||
from sglang.srt.layers.logits_processor import LogitsProcessorOutput, SamplingMaskStatus
|
||||
from sglang.srt.managers.schedule_batch import FINISH_ABORT, ReqKvInfo
|
||||
from sglang.srt.managers.scheduler_components.batch_result_processor import (
|
||||
SchedulerBatchResultProcessor,
|
||||
)
|
||||
from sglang.srt.managers.utils import GenerationBatchResult
|
||||
from sglang.srt.runtime_context import get_context
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=1, suite="base-a-test-cpu")
|
||||
@@ -205,5 +210,54 @@ def test_aborted_result_releases_mamba_allocated_before_kv():
|
||||
scheduler.output_streamer.stream_output.assert_called_once_with([req], False)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("transport_error", [False, True])
|
||||
@pytest.mark.parametrize(
|
||||
"status,http_status,err_type",
|
||||
[
|
||||
(SamplingMaskStatus.OVERFLOW, 400, "BadRequestError"),
|
||||
(SamplingMaskStatus.INVALID, 500, "InternalServerError"),
|
||||
],
|
||||
)
|
||||
@patch("sglang.srt.disaggregation.prefill.release_kv_cache", side_effect=_free_req)
|
||||
def test_sampling_mask_abort_preserves_error_and_releases_once(
|
||||
release_kv_cache, status, http_status, err_type, transport_error
|
||||
):
|
||||
"""A failed sender notification must not leak ownership or lose the API error."""
|
||||
scheduler = _Scheduler()
|
||||
scheduler.batch_result_processor.get_sampling_mask_finish_reason = lambda **kwargs: (
|
||||
SchedulerBatchResultProcessor.get_sampling_mask_finish_reason(None, **kwargs)
|
||||
)
|
||||
req = _Req(inflight_middle_chunks=0)
|
||||
req.to_finish = None
|
||||
req.return_sampling_mask = True
|
||||
req.time_stats.trace_ctx = Mock()
|
||||
if transport_error:
|
||||
req.disagg_kv_sender.abort.side_effect = RuntimeError("transport is down")
|
||||
result = GenerationBatchResult(
|
||||
next_token_ids=torch.tensor([11]),
|
||||
logits_output=LogitsProcessorOutput(
|
||||
next_token_logits=None, next_token_sampling_mask_status=[status]
|
||||
),
|
||||
)
|
||||
|
||||
with get_context().override_server_args(sampling_mask_max_tokens=64):
|
||||
scheduler.process_batch_result_disagg_prefill(_batch(req), result)
|
||||
scheduler.process_batch_result_disagg_prefill(_batch(req), result)
|
||||
|
||||
assert req.finished_reason.status_code == http_status
|
||||
assert req.finished_reason.err_type == err_type
|
||||
assert req.output_ids == []
|
||||
assert not req.kv.holds_kv and not req.kv.holds_mamba
|
||||
assert req.metadata_buffer_index == -1
|
||||
assert not req.pending_bootstrap
|
||||
assert req.rid not in scheduler.disagg_prefill_pending_chunk_rids
|
||||
release_kv_cache.assert_called_once_with(req, scheduler.tree_cache, is_insert=False)
|
||||
req.disagg_kv_sender.abort.assert_called_once_with()
|
||||
scheduler.req_to_metadata_buffer_idx_allocator.free.assert_called_once_with(7)
|
||||
scheduler.tree_cache.release_aborted_request.assert_called_once_with(req.rid)
|
||||
scheduler.output_streamer.stream_output.assert_called_once_with([req], False)
|
||||
scheduler.send_kv_chunk.assert_not_called()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(pytest.main([__file__, "-v"]))
|
||||
|
||||
@@ -9,6 +9,7 @@ Requires: torch, sglang (run in an environment with sglang installed)
|
||||
|
||||
import gc
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock
|
||||
from weakref import WeakKeyDictionary as WeakKeyDict
|
||||
|
||||
@@ -20,9 +21,14 @@ from sglang.srt.disaggregation.decode_kvcache_offload_manager import (
|
||||
from sglang.srt.disaggregation.kv_events import OffloadedState
|
||||
from sglang.srt.managers.cache_controller import HiCacheAck
|
||||
from sglang.srt.managers.schedule_batch import ReqKvInfo
|
||||
from sglang.srt.managers.scheduler_components.batch_result_processor import (
|
||||
SchedulerBatchResultProcessor,
|
||||
)
|
||||
from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
|
||||
from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
|
||||
from sglang.srt.runtime_context import get_context
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cpu_ci(est_time=10, suite="base-a-test-cpu")
|
||||
|
||||
@@ -440,5 +446,46 @@ class TestReleaseFinishedReq(unittest.TestCase):
|
||||
self.assertEqual(len(manager.offload_inflight), 0)
|
||||
|
||||
|
||||
class TestSamplingMaskAbortOffload(CustomTestCase):
|
||||
def test_abort_waits_for_existing_offload_before_reusing_slots(self):
|
||||
"""An abort must not recycle slots while a previous D2H copy reads them."""
|
||||
for inflight in (False, True):
|
||||
with self.subTest(inflight=inflight):
|
||||
manager, freed = _make_manager(pool_size=32)
|
||||
req = _make_mock_req(0, 20, 20)
|
||||
req.multimodal_inputs = None
|
||||
req.finished.return_value = True
|
||||
manager.req_to_token_pool.free.side_effect = lambda req: setattr(
|
||||
req.kv, "req_pool_idx", None
|
||||
)
|
||||
processor = SimpleNamespace(decode_offload_manager=manager)
|
||||
if inflight:
|
||||
manager.offload_inflight[req] = 1
|
||||
manager.ongoing_offload[1] = (req, torch.arange(4), [1], 0.0)
|
||||
manager.cache_controller = MagicMock()
|
||||
manager.cache_controller.ack_write_queue = [
|
||||
HiCacheAck(None, _FinishedEvent(), [1])
|
||||
]
|
||||
manager._trigger_backup = MagicMock(return_value="hash")
|
||||
|
||||
with get_context().override_server_args(
|
||||
disaggregation_decode_enable_offload_kvcache=True,
|
||||
enable_hisparse=False,
|
||||
):
|
||||
SchedulerBatchResultProcessor._handle_sampling_mask_abort(
|
||||
processor, req
|
||||
)
|
||||
|
||||
if inflight:
|
||||
self.assertEqual(freed, [])
|
||||
self.assertEqual(req.kv.req_pool_idx, 0)
|
||||
manager._check_offload_progress(1)
|
||||
self.assertEqual(len(freed), 1)
|
||||
self.assertTrue(torch.equal(freed[0], torch.arange(20)))
|
||||
self.assertIsNone(req.kv.req_pool_idx)
|
||||
manager.finalize_release_on_finish(req)
|
||||
self.assertEqual(len(freed), 1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -4,6 +4,8 @@ from unittest.mock import Mock, patch
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.layers.logits_processor import LogitsProcessorOutput, SamplingMaskStatus
|
||||
from sglang.srt.managers.scheduler_components.batch_result_processor import (
|
||||
SchedulerBatchResultProcessor,
|
||||
)
|
||||
@@ -49,6 +51,31 @@ def _make_processor(case, server_mode: str = "full") -> SchedulerBatchResultProc
|
||||
)
|
||||
|
||||
|
||||
class TestSamplingMaskMaterialization(CustomTestCase):
|
||||
def test_packed_ids_are_copied_before_per_request_slicing(self):
|
||||
"""Non-overlap capture must not perform one device copy per request."""
|
||||
packed_ids = Mock()
|
||||
packed_ids.shape = (2, 3)
|
||||
packed_ids.cpu.return_value = torch.tensor([[7, 8, 0], [9, 0, 0]])
|
||||
output = LogitsProcessorOutput(
|
||||
next_token_logits=None,
|
||||
sampling_mask_output=SimpleNamespace(
|
||||
token_ids=packed_ids,
|
||||
lengths=torch.tensor([2, 1]),
|
||||
selected_logprobs=torch.tensor([-0.5, -0.25]),
|
||||
statuses=torch.tensor([SamplingMaskStatus.OK, SamplingMaskStatus.OK]),
|
||||
),
|
||||
)
|
||||
SchedulerBatchResultProcessor.materialize_sampling_mask_output(
|
||||
reqs=[SimpleNamespace(return_sampling_mask=x) for x in (True, False, True)],
|
||||
output=output,
|
||||
)
|
||||
packed_ids.cpu.assert_called_once_with()
|
||||
self.assertEqual(output.next_token_sampling_mask_idx, [[7, 8], None, [9]])
|
||||
self.assertEqual(output.next_token_sampling_logprobs, [-0.5, None, -0.25])
|
||||
self.assertIsNone(output.sampling_mask_output)
|
||||
|
||||
|
||||
class _PrefillReq:
|
||||
def __init__(self, *, rid: str, inflight_middle_chunks: int, return_hidden_states):
|
||||
self.rid = rid
|
||||
@@ -138,6 +165,7 @@ class TestPrefillHiddenStateOffsets(CustomTestCase):
|
||||
logits_output=SimpleNamespace(
|
||||
hidden_states=hidden_states,
|
||||
customized_info=None,
|
||||
sampling_mask_output=None,
|
||||
),
|
||||
next_token_ids=torch.tensor([0, 1]),
|
||||
extend_input_len_per_req=[2, 3],
|
||||
@@ -165,6 +193,134 @@ class TestPrefillHiddenStateOffsets(CustomTestCase):
|
||||
self.assertEqual(last.hidden_states, [[22.0]])
|
||||
|
||||
|
||||
class TestPrefillSkippedOutput(CustomTestCase):
|
||||
def test_sampling_mask_middle_chunk_does_not_require_logits_output(self):
|
||||
"""A non-token-producing PP chunk may omit its logits output."""
|
||||
req = _PrefillReq(
|
||||
rid="middle",
|
||||
inflight_middle_chunks=1,
|
||||
return_hidden_states=False,
|
||||
)
|
||||
req.return_sampling_mask = True
|
||||
batch = SimpleNamespace(
|
||||
reqs=[req],
|
||||
return_logprob=False,
|
||||
return_hidden_states=False,
|
||||
return_hidden_states_mode=CaptureHiddenMode.NULL,
|
||||
spec_info=None,
|
||||
prefill_stats=None,
|
||||
dp_cooperation_info=None,
|
||||
)
|
||||
result = SimpleNamespace(
|
||||
copy_done=None,
|
||||
auxiliary_host_output=None,
|
||||
routed_experts_output=None,
|
||||
indexer_topk_output=None,
|
||||
logits_output=None,
|
||||
next_token_ids=torch.zeros(1, dtype=torch.int64),
|
||||
extend_input_len_per_req=None,
|
||||
extend_logprob_start_len_per_req=None,
|
||||
grammar_advanced=False,
|
||||
can_run_cuda_graph=False,
|
||||
skipped_output_comm=True,
|
||||
)
|
||||
processor = _make_processor(self)
|
||||
|
||||
with patch.object(
|
||||
envs.SGLANG_PP_SKIP_PURE_CHUNKED_OUTPUT_COMM,
|
||||
"get",
|
||||
return_value=True,
|
||||
):
|
||||
processor.process_batch_result_prefill(batch, result)
|
||||
|
||||
self.assertEqual(req.inflight_middle_chunks, 0)
|
||||
self.assertEqual(req.output_ids, [])
|
||||
processor.output_streamer.stream_output.assert_called_once_with(
|
||||
[req], False, req
|
||||
)
|
||||
|
||||
|
||||
class TestDecodeWithoutLogits(CustomTestCase):
|
||||
def test_pipeline_result_commits_token_without_sampling_metadata(self):
|
||||
processor = _make_processor(self)
|
||||
req = _DecodeReq()
|
||||
req.return_hidden_states = False
|
||||
batch = SimpleNamespace(
|
||||
reqs=[req],
|
||||
return_logprob=False,
|
||||
spec_algorithm=SimpleNamespace(is_none=lambda: True),
|
||||
batch_size=lambda: 1,
|
||||
)
|
||||
result = GenerationBatchResult(
|
||||
logits_output=None,
|
||||
next_token_ids=torch.tensor([8]),
|
||||
)
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
SchedulerBatchResultProcessor, "_maybe_update_reasoning_tokens"
|
||||
),
|
||||
patch.object(
|
||||
SchedulerBatchResultProcessor, "_handle_finish_state_updated_req"
|
||||
),
|
||||
):
|
||||
processor.process_batch_result_decode(batch, result)
|
||||
|
||||
self.assertEqual(req.output_ids, [8])
|
||||
self.assertEqual(processor.metrics_reporter.num_generated_tokens, 1)
|
||||
processor.output_streamer.stream_output.assert_called_once_with([req], False)
|
||||
|
||||
|
||||
class TestSamplingMaskStatusErrors(CustomTestCase):
|
||||
def test_decode_abort_releases_cache_without_committing_token(self):
|
||||
processor = _make_processor(self)
|
||||
req = _DecodeReq()
|
||||
req.output_ids = [7]
|
||||
req.return_sampling_mask = True
|
||||
req.multimodal_inputs = None
|
||||
req.update_finish_state = Mock()
|
||||
batch = SimpleNamespace(
|
||||
reqs=[req],
|
||||
return_logprob=False,
|
||||
spec_algorithm=SimpleNamespace(is_none=lambda: True),
|
||||
batch_size=lambda: 1,
|
||||
)
|
||||
result = GenerationBatchResult(
|
||||
logits_output=LogitsProcessorOutput(
|
||||
next_token_logits=None,
|
||||
next_token_sampling_mask_status=[SamplingMaskStatus.OVERFLOW],
|
||||
),
|
||||
next_token_ids=torch.tensor([8]),
|
||||
)
|
||||
with patch(
|
||||
"sglang.srt.managers.scheduler_components.batch_result_processor.release_kv_cache"
|
||||
) as release:
|
||||
processor.process_batch_result_decode(batch, result)
|
||||
|
||||
self.assertEqual(req.output_ids, [7])
|
||||
self.assertEqual(req.to_finish.status_code, 400)
|
||||
req.update_finish_state.assert_called_once_with(0)
|
||||
processor.model_worker.prepare_for_kv_cache_release.assert_called_once_with(req)
|
||||
release.assert_called_once_with(req, processor.tree_cache, is_insert=False)
|
||||
processor.output_streamer.stream_output.assert_called_once_with([req], False)
|
||||
|
||||
def test_overflow_and_invalid_have_distinct_http_errors(self):
|
||||
processor = _make_processor(self)
|
||||
|
||||
overflow = processor.get_sampling_mask_finish_reason(
|
||||
status=SamplingMaskStatus.OVERFLOW
|
||||
)
|
||||
self.assertEqual(overflow.status_code, 400)
|
||||
self.assertEqual(overflow.err_type, "BadRequestError")
|
||||
self.assertIn("cutoff ties", overflow.message)
|
||||
|
||||
invalid = processor.get_sampling_mask_finish_reason(
|
||||
status=SamplingMaskStatus.INVALID
|
||||
)
|
||||
self.assertEqual(invalid.status_code, 500)
|
||||
self.assertEqual(invalid.err_type, "InternalServerError")
|
||||
|
||||
|
||||
class TestDecodeHiddenStateRetention(CustomTestCase):
|
||||
def test_last_mode_multi_step_storage_stays_bounded(self):
|
||||
processor = _make_processor(self)
|
||||
@@ -180,7 +336,9 @@ class TestDecodeHiddenStateRetention(CustomTestCase):
|
||||
|
||||
def result(hidden_states):
|
||||
return GenerationBatchResult(
|
||||
logits_output=SimpleNamespace(hidden_states=hidden_states),
|
||||
logits_output=SimpleNamespace(
|
||||
hidden_states=hidden_states, sampling_mask_output=None
|
||||
),
|
||||
speculative_num_draft_tokens=4,
|
||||
)
|
||||
|
||||
|
||||
@@ -80,7 +80,9 @@ def _make_processor() -> SchedulerBatchResultProcessor:
|
||||
|
||||
def _make_result():
|
||||
return GenerationBatchResult(
|
||||
logits_output=SimpleNamespace(hidden_states=None, customized_info=None),
|
||||
logits_output=SimpleNamespace(
|
||||
hidden_states=None, customized_info=None, sampling_mask_output=None
|
||||
),
|
||||
next_token_ids=[4],
|
||||
speculative_num_draft_tokens=0,
|
||||
)
|
||||
|
||||
@@ -6,13 +6,17 @@ import pytest
|
||||
import torch
|
||||
|
||||
from sglang.srt.disaggregation.prefill import SchedulerDisaggregationPrefillMixin
|
||||
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
|
||||
from sglang.srt.layers.logits_processor import (
|
||||
LogitsProcessorOutput,
|
||||
SamplingMaskOutput,
|
||||
SamplingMaskStatus,
|
||||
)
|
||||
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
|
||||
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
|
||||
@@ -112,6 +116,58 @@ def test_auxiliary_output_releases_device_holder_after_copy():
|
||||
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(
|
||||
SimpleNamespace(), 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(
|
||||
|
||||
@@ -0,0 +1,76 @@
|
||||
"""Reject PD sampling-mask requests before entering transfer queues."""
|
||||
|
||||
import unittest
|
||||
from http import HTTPStatus
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel
|
||||
|
||||
maybe_stub_sgl_kernel()
|
||||
|
||||
from sglang.srt.disaggregation.utils import DisaggregationMode
|
||||
from sglang.srt.managers.scheduler import Scheduler
|
||||
|
||||
register_cpu_ci(est_time=2, suite="base-a-test-cpu")
|
||||
|
||||
|
||||
class TestSchedulerSamplingMaskValidation(CustomTestCase):
|
||||
def test_disabled_pd_masks_return_bad_request_before_admission(self):
|
||||
"""Neither PD role may queue mask requests without metadata buffers."""
|
||||
for mode in (DisaggregationMode.PREFILL, DisaggregationMode.DECODE):
|
||||
with self.subTest(mode=mode):
|
||||
scheduler = Scheduler.__new__(Scheduler)
|
||||
scheduler.enable_session_radix_cache = False
|
||||
scheduler.model_config = SimpleNamespace(
|
||||
hf_eos_token_id={1}, vocab_size=128
|
||||
)
|
||||
scheduler.disaggregation_mode = mode
|
||||
scheduler.disagg_metadata_buffers = SimpleNamespace(
|
||||
enable_sampling_mask=False
|
||||
)
|
||||
scheduler.metrics_reporter = SimpleNamespace(enable_metrics=False)
|
||||
scheduler.tokenizer = None
|
||||
scheduler.dllm_config = None
|
||||
scheduler._maybe_namespace_elastic_radix_cache = MagicMock()
|
||||
scheduler.spec_algorithm = SimpleNamespace(
|
||||
is_dflash_family=lambda: False,
|
||||
is_uno=lambda: False,
|
||||
)
|
||||
scheduler._add_request_to_queue = MagicMock()
|
||||
scheduler.output_streamer = MagicMock()
|
||||
recv_req = MagicMock(
|
||||
session_params=None,
|
||||
session_id=None,
|
||||
input_embeds=None,
|
||||
bootstrap_port=1,
|
||||
bootstrap_room=9,
|
||||
)
|
||||
req = MagicMock(return_sampling_mask=True, return_logprob=False)
|
||||
with (
|
||||
patch(
|
||||
"sglang.srt.managers.scheduler.BeamCoordinator.request_beam_width",
|
||||
return_value=1,
|
||||
),
|
||||
patch("sglang.srt.managers.scheduler.Req", return_value=req),
|
||||
patch("sglang.srt.managers.scheduler.prepare_abort") as abort,
|
||||
):
|
||||
scheduler.handle_generate_request(recv_req)
|
||||
abort.assert_called_once()
|
||||
self.assertIs(abort.call_args.args[0], req)
|
||||
self.assertIn(
|
||||
"SGLANG_ENABLE_DISAGG_SAMPLING_MASK=1",
|
||||
abort.call_args.args[1],
|
||||
)
|
||||
self.assertEqual(
|
||||
abort.call_args.kwargs["status_code"], HTTPStatus.BAD_REQUEST
|
||||
)
|
||||
scheduler.output_streamer.stream_output.assert_called_once_with(
|
||||
[req], False
|
||||
)
|
||||
scheduler._add_request_to_queue.assert_not_called()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -100,6 +100,54 @@ class TestSamplingBatchInfoLen(CustomTestCase):
|
||||
self.assertEqual(len(info), 5)
|
||||
|
||||
|
||||
class TestSamplingMaskBatchIndices(CustomTestCase):
|
||||
def test_filter_removes_last_opted_in_row_then_merge_restores_capture(self):
|
||||
info = _make_info(
|
||||
batch_size=2,
|
||||
return_sampling_masks=[False, True],
|
||||
sampling_mask_batch_indices=torch.tensor([1]),
|
||||
)
|
||||
info.filter_batch([0], torch.tensor([0]))
|
||||
self.assertIsNone(info.sampling_mask_batch_indices)
|
||||
other = _make_info(
|
||||
batch_size=1,
|
||||
return_sampling_masks=[True],
|
||||
sampling_mask_batch_indices=torch.tensor([0]),
|
||||
)
|
||||
info.merge_batch(other)
|
||||
self.assertEqual(info.return_sampling_masks, [False, True])
|
||||
self.assertEqual(info.sampling_mask_batch_indices.tolist(), [1])
|
||||
|
||||
def test_filter_rebuilds_row_indices(self):
|
||||
info = _make_info(
|
||||
batch_size=4,
|
||||
return_sampling_masks=[False, True, False, True],
|
||||
sampling_mask_batch_indices=torch.tensor([1, 3]),
|
||||
)
|
||||
|
||||
info.filter_batch([1, 2, 3], torch.tensor([1, 2, 3]))
|
||||
|
||||
self.assertEqual(info.return_sampling_masks, [True, False, True])
|
||||
self.assertEqual(info.sampling_mask_batch_indices.tolist(), [0, 2])
|
||||
|
||||
def test_merge_offsets_rhs_row_indices(self):
|
||||
lhs = _make_info(
|
||||
batch_size=2,
|
||||
return_sampling_masks=[False, True],
|
||||
sampling_mask_batch_indices=torch.tensor([1]),
|
||||
)
|
||||
rhs = _make_info(
|
||||
batch_size=3,
|
||||
return_sampling_masks=[True, False, True],
|
||||
sampling_mask_batch_indices=torch.tensor([0, 2]),
|
||||
)
|
||||
|
||||
lhs.merge_batch(rhs)
|
||||
|
||||
self.assertEqual(lhs.return_sampling_masks, [False, True, True, False, True])
|
||||
self.assertEqual(lhs.sampling_mask_batch_indices.tolist(), [1, 2, 4])
|
||||
|
||||
|
||||
class TestMergeCustomLogitProcessor(CustomTestCase):
|
||||
def test_both_none_returns_none(self):
|
||||
"""Test that merging two None processor dicts returns None."""
|
||||
|
||||
@@ -186,6 +186,38 @@ class TestPrepareServerArgs(CustomTestCase):
|
||||
):
|
||||
ServerArgs(model_path="dummy", prefill_decode_interval=-1).resolve_once()
|
||||
|
||||
def test_sampling_mask_max_tokens(self):
|
||||
self.assertEqual(ServerArgs(model_path="dummy").sampling_mask_max_tokens, 4096)
|
||||
self.assertEqual(
|
||||
ServerArgs(
|
||||
model_path="dummy", sampling_mask_max_tokens=8192
|
||||
).sampling_mask_max_tokens,
|
||||
8192,
|
||||
)
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, "--sampling-mask-max-tokens must be positive"
|
||||
):
|
||||
prepare_server_args(
|
||||
["--model-path", "dummy", "--sampling-mask-max-tokens", "0"]
|
||||
).resolve_once()
|
||||
|
||||
def test_legacy_sampling_mask_env_requires_migration(self):
|
||||
"""Legacy configuration must not silently disable PD sampling masks."""
|
||||
for value in ("0", "128", "invalid", ""):
|
||||
for enabled in (False, True):
|
||||
with (
|
||||
self.subTest(value=value, enabled=enabled),
|
||||
envs.SGLANG_DISAGGREGATION_SAMPLING_MASK_MAX_TOKENS.override(value),
|
||||
envs.SGLANG_ENABLE_DISAGG_SAMPLING_MASK.override(enabled),
|
||||
self.assertRaisesRegex(
|
||||
ValueError,
|
||||
"SGLANG_DISAGGREGATION_SAMPLING_MASK_MAX_TOKENS.*"
|
||||
"Unset it.*SGLANG_ENABLE_DISAGG_SAMPLING_MASK=1.*"
|
||||
"--sampling-mask-max-tokens.*prefill and decode",
|
||||
),
|
||||
):
|
||||
ServerArgs(model_path="dummy").resolve_once()
|
||||
|
||||
def test_dsv4_prefill_backend_cli_choices(self):
|
||||
parser = server_args_module.argparse.ArgumentParser()
|
||||
ServerArgs.add_cli_args(parser)
|
||||
|
||||
Reference in New Issue
Block a user