Add a deterministic token oracle and production write-input assertion (#26815)
This commit is contained in:
fzyzcjy
@@ -0,0 +1,86 @@
|
||||
"""Regression for PR #25015 EAGLE positions misalign: revert the fix and expect canary fire."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import unittest
|
||||
from typing import ClassVar
|
||||
|
||||
from sglang.srt.kv_canary.config import CanaryMode
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.kv_canary.e2e_base import CanaryE2EBase
|
||||
|
||||
register_cuda_ci(est_time=60, stage="extra-a", runner_config="1-gpu-small")
|
||||
|
||||
_SPEC_EAGLE_TOKEN_ORACLE_ENV = {
|
||||
"SGLANG_KV_CANARY_ENABLE_WRITE_INPUT_ASSERT": "0",
|
||||
"SGLANG_KV_CANARY_ENABLE_TOKEN_ORACLE": "1",
|
||||
}
|
||||
_SPEC_EAGLE_REVERT_PR_ENV = {
|
||||
**_SPEC_EAGLE_TOKEN_ORACLE_ENV,
|
||||
"SGLANG_DEBUG_REVERT_PR": "25015",
|
||||
}
|
||||
_CUDA_GRAPH_MAX_BS = 1
|
||||
_EAGER_DRAFT_REQUEST_COUNT = 20
|
||||
assert _EAGER_DRAFT_REQUEST_COUNT > _CUDA_GRAPH_MAX_BS
|
||||
|
||||
_SPEC_EAGLE_SERVER_ARGS = (
|
||||
"--sampling-backend",
|
||||
"token_oracle",
|
||||
"--speculative-algorithm",
|
||||
"EAGLE",
|
||||
"--cuda-graph-max-bs",
|
||||
str(_CUDA_GRAPH_MAX_BS),
|
||||
"--max-running-requests",
|
||||
"32",
|
||||
)
|
||||
|
||||
|
||||
class _EaglePositionsBase(CanaryE2EBase):
|
||||
model_mode = "mha"
|
||||
# LOG mode keeps the server alive after the first violation so server warmup + this test's
|
||||
# parallel requests both run; we then read the violation log to assert the position bit fired.
|
||||
kv_canary_mode = CanaryMode.LOG
|
||||
extra_server_args = _SPEC_EAGLE_SERVER_ARGS
|
||||
revert_pr: ClassVar[bool]
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls) -> None:
|
||||
if cls is _EaglePositionsBase:
|
||||
raise unittest.SkipTest("abstract base; concrete subclasses set revert_pr")
|
||||
cls.extra_env = (
|
||||
_SPEC_EAGLE_REVERT_PR_ENV if cls.revert_pr else _SPEC_EAGLE_TOKEN_ORACLE_ENV
|
||||
)
|
||||
super().setUpClass()
|
||||
|
||||
def test_pr_25015_eagle_positions(self) -> None:
|
||||
self.send_parallel_requests(
|
||||
n=_EAGER_DRAFT_REQUEST_COUNT,
|
||||
assert_all_success=not self.revert_pr,
|
||||
max_new_tokens=32,
|
||||
timeout=60.0,
|
||||
)
|
||||
|
||||
if self.revert_pr:
|
||||
self.assert_violation_logged_any(
|
||||
launch_tag_patterns=("*",),
|
||||
fail_reason="verify_position",
|
||||
flush_wait_seconds=0.0,
|
||||
)
|
||||
else:
|
||||
self.assert_no_violation(wait_seconds=2.0)
|
||||
|
||||
|
||||
class TestEaglePositionsMisalignRegression(_EaglePositionsBase):
|
||||
"""Revert PR #25015 fix and expect canary to fire a position-mismatch violation."""
|
||||
|
||||
revert_pr = True
|
||||
|
||||
|
||||
class TestEaglePositionsMatchWithFix(_EaglePositionsBase):
|
||||
"""With the PR #25015 fix in place, no canary fires."""
|
||||
|
||||
revert_pr = False
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,57 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.kv_canary.expected_inputs import ExpectedInputs
|
||||
from sglang.srt.kv_canary.token_oracle.oracle import HashOracle
|
||||
from sglang.srt.kv_canary.token_oracle.oracle_manager import TokenOracleManager
|
||||
from sglang.srt.model_executor.forward_batch_info import ForwardMode
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.kv_canary.fixtures import DEFAULT_DEVICE
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cuda_ci(est_time=1, stage="extra-a", runner_config="1-gpu-small")
|
||||
|
||||
|
||||
class TestTokenOracleManager(CustomTestCase):
|
||||
def setUp(self) -> None:
|
||||
self.device = DEFAULT_DEVICE
|
||||
|
||||
def test_fill_expected_inputs_expands_draft_extend_generalized_req_ids_per_token(
|
||||
self,
|
||||
) -> None:
|
||||
"""Verify EAGLE draft extend maps one request row to every draft token."""
|
||||
forward_batch = SimpleNamespace(
|
||||
forward_mode=ForwardMode.DRAFT_EXTEND,
|
||||
spec_info=SimpleNamespace(num_tokens_per_req=4),
|
||||
rids_int=torch.tensor([3, 7], dtype=torch.int64, device=self.device),
|
||||
bootstrap_room_ids_int=None,
|
||||
input_ids=torch.tensor(
|
||||
[101, 102, 103, 104, 201, 202, 203, 204],
|
||||
dtype=torch.int64,
|
||||
device=self.device,
|
||||
),
|
||||
positions=torch.arange(8, dtype=torch.int64, device=self.device),
|
||||
extend_seq_lens=torch.tensor([1, 1], dtype=torch.int64, device=self.device),
|
||||
)
|
||||
expected_inputs = ExpectedInputs.allocate(capacity=8, device=self.device)
|
||||
manager = TokenOracleManager(oracle=HashOracle(vocab_size=32000))
|
||||
|
||||
manager.fill_expected_inputs(
|
||||
forward_batch=forward_batch,
|
||||
expected_inputs_out=expected_inputs,
|
||||
)
|
||||
|
||||
self.assertTrue(
|
||||
torch.equal(expected_inputs.tokens[:8], forward_batch.input_ids)
|
||||
)
|
||||
self.assertTrue(
|
||||
torch.equal(expected_inputs.positions[:8], forward_batch.positions)
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,36 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import unittest
|
||||
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.mock_model.perturb_e2e_base import MockModelPerturbE2EBase
|
||||
|
||||
register_cuda_ci(est_time=60, stage="extra-a", runner_config="1-gpu-small")
|
||||
|
||||
|
||||
class TestPerturbNextTokenSwap(MockModelPerturbE2EBase):
|
||||
"""Mock-model self-test: swap two reqs' sampled next tokens at the sampler exit.
|
||||
|
||||
KV path is untouched, so kv_canary KV-side fail_reasons stay silent. The
|
||||
token-oracle input check downstream MUST report fail_reason=write_token.
|
||||
Validates the input-check link is genuinely active.
|
||||
"""
|
||||
|
||||
extra_env = {
|
||||
"SGLANG_KV_CANARY_PERTURB_NEXT_TOKEN_SWAP_PROB": "0.1",
|
||||
"SGLANG_KV_CANARY_PERTURB_WARMUP_STEPS": "0",
|
||||
}
|
||||
extra_server_args = ("--skip-server-warmup",)
|
||||
|
||||
def test_swap_triggers_input_check_violation_but_kv_paths_silent(self) -> None:
|
||||
"""Verify next_token swap fires write_token violation while KV reasons stay silent."""
|
||||
self.send_parallel_requests(n=4, timeout=30.0)
|
||||
self.assert_log_contains("kv_canary perturb next_token_swap: swapped")
|
||||
self.assert_any_launch_tag_violation_reported(fail_reason="write_token")
|
||||
self.assert_any_launch_tag_violation_absent(fail_reason="verify_real_kv_hash")
|
||||
self.assert_any_launch_tag_violation_absent(fail_reason="verify_position")
|
||||
self.assert_any_launch_tag_violation_absent(fail_reason="verify_chain_hash")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,238 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import dataclasses
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.kv_canary.expected_inputs import ExpectedInputs
|
||||
from sglang.srt.kv_canary.token_oracle.oracle import HashOracle
|
||||
from sglang.srt.kv_canary.token_oracle.sampler import install_oracle_sampler
|
||||
from sglang.srt.model_executor.forward_batch_info import (
|
||||
ForwardMode,
|
||||
_stable_hash_str_to_i64,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.mock_model.utils import mock_model_server_args, mock_model_server_env
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cuda_ci(est_time=60, stage="extra-a", runner_config="1-gpu-small")
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class _StubForwardBatch:
|
||||
input_ids: torch.Tensor
|
||||
positions: torch.Tensor
|
||||
req_pool_indices: torch.Tensor
|
||||
forward_mode: ForwardMode
|
||||
extend_seq_lens: object
|
||||
rids_int: torch.Tensor
|
||||
bootstrap_room_ids_int: torch.Tensor | None = None
|
||||
spec_info: object | None = None
|
||||
seq_lens: torch.Tensor | None = None
|
||||
|
||||
|
||||
def _scalar_expected_token(
|
||||
oracle: HashOracle, *, generalized_req_id: int, position: int
|
||||
) -> int:
|
||||
out = oracle.expected_tokens(
|
||||
generalized_req_ids=torch.tensor([generalized_req_id], dtype=torch.int64),
|
||||
positions=torch.tensor([position], dtype=torch.int64),
|
||||
)
|
||||
return int(out.tolist()[0])
|
||||
|
||||
|
||||
class TestFillExpectedInputs(CustomTestCase):
|
||||
def test_sample_next_tokens_uses_next_position(self) -> None:
|
||||
oracle = HashOracle(vocab_size=32000)
|
||||
hook = install_oracle_sampler(oracle=oracle)
|
||||
|
||||
rid_a = "req-a"
|
||||
hashed_a = _stable_hash_str_to_i64(rid_a)
|
||||
out = hook.sample_next_tokens(
|
||||
generalized_req_ids=torch.tensor([hashed_a], dtype=torch.int64),
|
||||
logits_positions=torch.tensor([5], dtype=torch.int64),
|
||||
)
|
||||
|
||||
self.assertEqual(
|
||||
out.tolist(),
|
||||
[_scalar_expected_token(oracle, generalized_req_id=hashed_a, position=6)],
|
||||
)
|
||||
|
||||
def test_fill_expected_inputs_decode_one_token_per_req(self) -> None:
|
||||
"""Verify decode mode fills one expected token per request."""
|
||||
oracle = HashOracle(vocab_size=32000)
|
||||
hook = install_oracle_sampler(oracle=oracle)
|
||||
|
||||
rid_a = "req-a"
|
||||
rid_b = "req-b"
|
||||
fb = _StubForwardBatch(
|
||||
input_ids=torch.tensor([0, 0], dtype=torch.int64),
|
||||
positions=torch.tensor([10, 20], dtype=torch.int64),
|
||||
req_pool_indices=torch.tensor([5, 7], dtype=torch.int64),
|
||||
forward_mode=ForwardMode.DECODE,
|
||||
extend_seq_lens=None,
|
||||
rids_int=torch.tensor(
|
||||
[_stable_hash_str_to_i64(rid_a), _stable_hash_str_to_i64(rid_b)],
|
||||
dtype=torch.int64,
|
||||
),
|
||||
)
|
||||
expected_inputs = ExpectedInputs.allocate(
|
||||
capacity=8, device=torch.device("cpu")
|
||||
)
|
||||
|
||||
hook.fill_expected_inputs(
|
||||
forward_batch=fb,
|
||||
expected_inputs_out=expected_inputs,
|
||||
)
|
||||
|
||||
self.assertEqual(
|
||||
expected_inputs.tokens[:2].tolist(),
|
||||
[
|
||||
_scalar_expected_token(
|
||||
oracle,
|
||||
generalized_req_id=_stable_hash_str_to_i64(rid_a),
|
||||
position=10,
|
||||
),
|
||||
_scalar_expected_token(
|
||||
oracle,
|
||||
generalized_req_id=_stable_hash_str_to_i64(rid_b),
|
||||
position=20,
|
||||
),
|
||||
],
|
||||
)
|
||||
self.assertEqual(expected_inputs.positions[:2].tolist(), [10, 20])
|
||||
|
||||
def test_fill_expected_inputs_prefers_bootstrap_room_ids(self) -> None:
|
||||
"""Verify PD oracle checks can key by bootstrap room without rewriting rids_int."""
|
||||
oracle = HashOracle(vocab_size=32000)
|
||||
hook = install_oracle_sampler(oracle=oracle)
|
||||
|
||||
rid_a = "prefill-local-rid"
|
||||
rid_b = "regular-rid"
|
||||
hashed_a = _stable_hash_str_to_i64(rid_a)
|
||||
hashed_b = _stable_hash_str_to_i64(rid_b)
|
||||
fb = _StubForwardBatch(
|
||||
input_ids=torch.tensor([0, 0], dtype=torch.int64),
|
||||
positions=torch.tensor([10, 20], dtype=torch.int64),
|
||||
req_pool_indices=torch.tensor([5, 7], dtype=torch.int64),
|
||||
forward_mode=ForwardMode.DECODE,
|
||||
extend_seq_lens=None,
|
||||
rids_int=torch.tensor([hashed_a, hashed_b], dtype=torch.int64),
|
||||
bootstrap_room_ids_int=torch.tensor([1234, -1], dtype=torch.int64),
|
||||
)
|
||||
expected_inputs = ExpectedInputs.allocate(
|
||||
capacity=8, device=torch.device("cpu")
|
||||
)
|
||||
|
||||
hook.fill_expected_inputs(
|
||||
forward_batch=fb,
|
||||
expected_inputs_out=expected_inputs,
|
||||
)
|
||||
|
||||
self.assertEqual(fb.rids_int.tolist(), [hashed_a, hashed_b])
|
||||
self.assertEqual(
|
||||
expected_inputs.tokens[:2].tolist(),
|
||||
[
|
||||
_scalar_expected_token(oracle, generalized_req_id=1234, position=10),
|
||||
_scalar_expected_token(
|
||||
oracle, generalized_req_id=hashed_b, position=20
|
||||
),
|
||||
],
|
||||
)
|
||||
self.assertEqual(expected_inputs.positions[:2].tolist(), [10, 20])
|
||||
|
||||
def test_fill_expected_inputs_extend_uses_forward_input_ids(self) -> None:
|
||||
"""Verify extend mode checks prompt tokens already present in the forward batch."""
|
||||
oracle = HashOracle(vocab_size=32000)
|
||||
hook = install_oracle_sampler(oracle=oracle)
|
||||
|
||||
rid_a = "req-a"
|
||||
rid_b = "req-b"
|
||||
hashed_a = _stable_hash_str_to_i64(rid_a)
|
||||
hashed_b = _stable_hash_str_to_i64(rid_b)
|
||||
fb = _StubForwardBatch(
|
||||
input_ids=torch.tensor([101, 102, 103, 201], dtype=torch.int64),
|
||||
positions=torch.tensor([0, 1, 2, 0], dtype=torch.int64),
|
||||
req_pool_indices=torch.tensor([5, 7], dtype=torch.int64),
|
||||
forward_mode=ForwardMode.EXTEND,
|
||||
extend_seq_lens=torch.tensor([3, 1], dtype=torch.int64),
|
||||
rids_int=torch.tensor([hashed_a, hashed_b], dtype=torch.int64),
|
||||
)
|
||||
expected_inputs = ExpectedInputs.allocate(
|
||||
capacity=8, device=torch.device("cpu")
|
||||
)
|
||||
|
||||
hook.fill_expected_inputs(
|
||||
forward_batch=fb,
|
||||
expected_inputs_out=expected_inputs,
|
||||
)
|
||||
|
||||
self.assertEqual(
|
||||
expected_inputs.tokens[:4].tolist(),
|
||||
[101, 102, 103, 201],
|
||||
)
|
||||
self.assertEqual(expected_inputs.positions[:4].tolist(), [0, 1, 2, 0])
|
||||
|
||||
def test_fill_expected_inputs_zero_tokens_is_noop(
|
||||
self,
|
||||
) -> None:
|
||||
"""Verify filling zero expected tokens leaves the output buffer unchanged."""
|
||||
hook = install_oracle_sampler(oracle=HashOracle(vocab_size=100))
|
||||
|
||||
rid_a = "req-a"
|
||||
rid_b = "req-b"
|
||||
fb = _StubForwardBatch(
|
||||
input_ids=torch.empty(0, dtype=torch.int64),
|
||||
positions=torch.empty(0, dtype=torch.int64),
|
||||
req_pool_indices=torch.tensor([5, 7], dtype=torch.int64),
|
||||
forward_mode=ForwardMode.DECODE,
|
||||
extend_seq_lens=None,
|
||||
rids_int=torch.tensor(
|
||||
[_stable_hash_str_to_i64(rid_a), _stable_hash_str_to_i64(rid_b)],
|
||||
dtype=torch.int64,
|
||||
),
|
||||
)
|
||||
expected_inputs = ExpectedInputs.allocate(
|
||||
capacity=4, device=torch.device("cpu")
|
||||
)
|
||||
initial_tokens = expected_inputs.tokens.clone()
|
||||
|
||||
hook.fill_expected_inputs(
|
||||
forward_batch=fb,
|
||||
expected_inputs_out=expected_inputs,
|
||||
)
|
||||
|
||||
self.assertEqual(expected_inputs.tokens.tolist(), initial_tokens.tolist())
|
||||
|
||||
|
||||
class TestMockModelServerLaunchHelpers(CustomTestCase):
|
||||
def test_mock_model_server_args_adds_canary_defaults(self) -> None:
|
||||
"""Verify mock model launch args include KV canary defaults before user args."""
|
||||
args = mock_model_server_args("--tp", "2")
|
||||
|
||||
self.assertIn("--load-format", args)
|
||||
self.assertIn("dummy", args)
|
||||
self.assertIn("--sampling-backend", args)
|
||||
self.assertIn("token_oracle", args)
|
||||
self.assertIn("--kv-canary", args)
|
||||
self.assertIn("raise", args)
|
||||
self.assertEqual(args[-2:], ["--tp", "2"])
|
||||
|
||||
def test_mock_model_server_env_enables_input_check_by_default(self) -> None:
|
||||
"""Verify mock model launch env enables canary input checking by default."""
|
||||
env = mock_model_server_env()
|
||||
|
||||
self.assertEqual(env["SGLANG_KV_CANARY_ENABLE_WRITE_INPUT_ASSERT"], "1")
|
||||
self.assertEqual(env["SGLANG_KV_CANARY_ENABLE_TOKEN_ORACLE"], "1")
|
||||
|
||||
def test_mock_model_server_env_can_disable_input_check(self) -> None:
|
||||
"""Verify mock model launch env can disable canary input checking."""
|
||||
env = mock_model_server_env(input_check_enabled=False)
|
||||
|
||||
self.assertEqual(env["SGLANG_KV_CANARY_ENABLE_WRITE_INPUT_ASSERT"], "0")
|
||||
self.assertEqual(env["SGLANG_KV_CANARY_ENABLE_TOKEN_ORACLE"], "1")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,50 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
os.environ["SGLANG_KV_CANARY_ENABLE_TOKEN_ORACLE"] = "1"
|
||||
|
||||
from sglang.srt.kv_canary.token_oracle.install import install_token_oracle_from_env
|
||||
from sglang.srt.kv_canary.token_oracle.oracle import HashOracle
|
||||
from sglang.srt.layers.sampler import _CUSTOM_SAMPLER_FACTORIES
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cuda_ci(est_time=60, stage="extra-a", runner_config="1-gpu-small")
|
||||
|
||||
|
||||
def _make_server_args(*, sampling_backend: str) -> SimpleNamespace:
|
||||
return SimpleNamespace(sampling_backend=sampling_backend)
|
||||
|
||||
|
||||
class TestInstallTokenOracleFromEnv(CustomTestCase):
|
||||
def test_install_token_oracle_from_env_disabled_returns_none(self) -> None:
|
||||
"""Verify server-arg-disabled token oracle installation (sampling_backend != 'token_oracle') returns no TokenOracleManager."""
|
||||
server_args = _make_server_args(sampling_backend="auto")
|
||||
hook = install_token_oracle_from_env(server_args=server_args, vocab_size=1000)
|
||||
self.assertIsNone(hook)
|
||||
|
||||
def test_install_token_oracle_from_env_enabled_registers_oracle_backend(
|
||||
self,
|
||||
) -> None:
|
||||
"""Verify token oracle installation via sampling_backend='token_oracle' registers the oracle backend."""
|
||||
server_args = _make_server_args(sampling_backend="token_oracle")
|
||||
hook = install_token_oracle_from_env(server_args=server_args, vocab_size=512)
|
||||
self.assertIsNotNone(hook)
|
||||
self.assertIn("token_oracle", _CUSTOM_SAMPLER_FACTORIES)
|
||||
|
||||
def test_install_token_oracle_from_env_enabled_returns_hook_with_hash_oracle(
|
||||
self,
|
||||
) -> None:
|
||||
"""Verify token oracle installation via sampling_backend='token_oracle' returns a TokenOracleManager wrapping a HashOracle."""
|
||||
server_args = _make_server_args(sampling_backend="token_oracle")
|
||||
hook = install_token_oracle_from_env(server_args=server_args, vocab_size=256)
|
||||
self.assertIsNotNone(hook)
|
||||
self.assertIsInstance(hook.oracle, HashOracle)
|
||||
self.assertEqual(hook.oracle.vocab_size, 256)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,127 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import random
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.kv_canary.consts import splitmix64
|
||||
from sglang.srt.kv_canary.token_oracle.oracle import (
|
||||
HashOracle,
|
||||
_splitmix64_tensor,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cuda_ci(est_time=60, stage="extra-a", runner_config="1-gpu-small")
|
||||
|
||||
|
||||
_U64_MASK: int = (1 << 64) - 1
|
||||
|
||||
|
||||
def _signed_to_unsigned_i64(value: int) -> int:
|
||||
return value & _U64_MASK
|
||||
|
||||
|
||||
def _call(oracle: HashOracle, *, generalized_req_id: int, position: int) -> int:
|
||||
out = oracle.expected_tokens(
|
||||
generalized_req_ids=torch.tensor([generalized_req_id], dtype=torch.int64),
|
||||
positions=torch.tensor([position], dtype=torch.int64),
|
||||
)
|
||||
return int(out.tolist()[0])
|
||||
|
||||
|
||||
class TestHashOracle(CustomTestCase):
|
||||
def test_hash_oracle_is_deterministic_for_same_inputs(self) -> None:
|
||||
"""Verify HashOracle returns the same token for identical inputs."""
|
||||
oracle = HashOracle(vocab_size=32000)
|
||||
|
||||
first = _call(oracle, generalized_req_id=7, position=42)
|
||||
second = _call(oracle, generalized_req_id=7, position=42)
|
||||
|
||||
self.assertEqual(first, second)
|
||||
|
||||
def test_hash_oracle_output_in_vocab_range(self) -> None:
|
||||
"""Verify HashOracle outputs stay within the configured vocabulary range."""
|
||||
vocab_size = 1024
|
||||
oracle = HashOracle(vocab_size=vocab_size)
|
||||
|
||||
generalized_req_ids = torch.arange(0, 64, dtype=torch.int64).repeat_interleave(
|
||||
64
|
||||
)
|
||||
positions = torch.arange(0, 64, dtype=torch.int64).repeat(64)
|
||||
tokens = oracle.expected_tokens(
|
||||
generalized_req_ids=generalized_req_ids, positions=positions
|
||||
).tolist()
|
||||
|
||||
for token in tokens:
|
||||
self.assertTrue(0 <= token < vocab_size)
|
||||
|
||||
|
||||
class TestSplitmix64Tensor(CustomTestCase):
|
||||
def test_splitmix64_tensor_matches_scalar_ref_on_random_inputs(self) -> None:
|
||||
"""Verify tensor SplitMix64 matches the scalar reference on random inputs."""
|
||||
rng = random.Random(0)
|
||||
num_cases = 1000
|
||||
unsigned_inputs: list[int] = [
|
||||
rng.randrange(0, 1 << 64) for _ in range(num_cases)
|
||||
]
|
||||
|
||||
signed_inputs = [
|
||||
value if value < (1 << 63) else value - (1 << 64)
|
||||
for value in unsigned_inputs
|
||||
]
|
||||
actual = _splitmix64_tensor(torch.tensor(signed_inputs, dtype=torch.int64))
|
||||
|
||||
actual_unsigned = [_signed_to_unsigned_i64(v) for v in actual.tolist()]
|
||||
expected_unsigned = [splitmix64(v) for v in unsigned_inputs]
|
||||
|
||||
self.assertEqual(actual_unsigned, expected_unsigned)
|
||||
|
||||
def test_splitmix64_tensor_known_vectors(self) -> None:
|
||||
"""Verify tensor SplitMix64 matches scalar reference values for known inputs."""
|
||||
inputs = [0, 1, -1, 1 << 32, (1 << 63) - 1, -(1 << 63)]
|
||||
expected_unsigned = [splitmix64(_signed_to_unsigned_i64(v)) for v in inputs]
|
||||
|
||||
actual = _splitmix64_tensor(torch.tensor(inputs, dtype=torch.int64))
|
||||
actual_unsigned = [_signed_to_unsigned_i64(v) for v in actual.tolist()]
|
||||
|
||||
self.assertEqual(actual_unsigned, expected_unsigned)
|
||||
|
||||
def test_splitmix64_tensor_preserves_shape_and_dtype(self) -> None:
|
||||
"""Verify tensor SplitMix64 preserves input shape and int64 dtype."""
|
||||
shape = (3, 4, 5)
|
||||
rng = torch.Generator().manual_seed(42)
|
||||
inputs = torch.randint(
|
||||
low=-(1 << 62),
|
||||
high=(1 << 62),
|
||||
size=shape,
|
||||
dtype=torch.int64,
|
||||
generator=rng,
|
||||
)
|
||||
|
||||
out = _splitmix64_tensor(inputs)
|
||||
|
||||
self.assertEqual(out.shape, inputs.shape)
|
||||
self.assertEqual(out.dtype, torch.int64)
|
||||
|
||||
def test_splitmix64_tensor_is_deterministic(self) -> None:
|
||||
"""Verify tensor SplitMix64 returns stable values for repeated calls."""
|
||||
inputs = torch.tensor([0, 1, 2, 3, 1 << 40, -7], dtype=torch.int64)
|
||||
|
||||
first = _splitmix64_tensor(inputs.clone()).tolist()
|
||||
second = _splitmix64_tensor(inputs.clone()).tolist()
|
||||
|
||||
self.assertEqual(first, second)
|
||||
|
||||
def test_splitmix64_tensor_is_injective_on_distinct_inputs(self) -> None:
|
||||
"""Verify tensor SplitMix64 maps distinct sampled inputs to distinct outputs."""
|
||||
inputs = torch.arange(-1000, 1000, dtype=torch.int64)
|
||||
|
||||
out = _splitmix64_tensor(inputs).tolist()
|
||||
|
||||
self.assertEqual(len(set(out)), len(out))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,85 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import random
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.kv_canary.verify_ref import splitmix64
|
||||
from sglang.srt.kv_canary.token_oracle.oracle import HashOracle
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cuda_ci(est_time=30, stage="extra-a", runner_config="1-gpu-small")
|
||||
|
||||
|
||||
class TestHashOracleTorchVsRef(CustomTestCase):
|
||||
def test_hash_oracle_matches_scalar_splitmix64_ref(self) -> None:
|
||||
"""Verify single-item HashOracle calls match the scalar SplitMix64 reference."""
|
||||
rng = random.Random(0)
|
||||
vocab_size = 32000
|
||||
num_cases = 1000
|
||||
|
||||
generalized_req_ids: list[int] = []
|
||||
positions: list[int] = []
|
||||
for _ in range(num_cases):
|
||||
generalized_req_ids.append(rng.randrange(0, 1 << 60))
|
||||
positions.append(rng.randrange(0, 1 << 60))
|
||||
|
||||
ref_tokens: list[int] = [
|
||||
splitmix64(generalized_req_ids[i] ^ positions[i]) % vocab_size
|
||||
for i in range(num_cases)
|
||||
]
|
||||
|
||||
oracle = HashOracle(vocab_size=vocab_size)
|
||||
torch_tokens: list[int] = []
|
||||
generalized_req_ids_tensor = torch.tensor(
|
||||
generalized_req_ids, dtype=torch.int64
|
||||
)
|
||||
positions_tensor = torch.tensor(positions, dtype=torch.int64)
|
||||
for i in range(num_cases):
|
||||
out = oracle.expected_tokens(
|
||||
generalized_req_ids=generalized_req_ids_tensor[i : i + 1],
|
||||
positions=positions_tensor[i : i + 1],
|
||||
)
|
||||
torch_tokens.append(int(out.tolist()[0]))
|
||||
|
||||
for i in range(num_cases):
|
||||
self.assertEqual(
|
||||
torch_tokens[i],
|
||||
ref_tokens[i],
|
||||
f"mismatch at case {i}: generalized_req_id={generalized_req_ids[i]} "
|
||||
f"position={positions[i]}: torch={torch_tokens[i]} ref={ref_tokens[i]}",
|
||||
)
|
||||
|
||||
def test_hash_oracle_batched_matches_scalar_splitmix64_ref(self) -> None:
|
||||
"""Verify batched HashOracle calls match the scalar SplitMix64 reference."""
|
||||
rng = random.Random(1)
|
||||
vocab_size = 32000
|
||||
num_cases = 1000
|
||||
|
||||
generalized_req_ids = [rng.randrange(0, 1 << 60) for _ in range(num_cases)]
|
||||
positions = [rng.randrange(0, 1 << 60) for _ in range(num_cases)]
|
||||
ref_tokens = [
|
||||
splitmix64(generalized_req_ids[i] ^ positions[i]) % vocab_size
|
||||
for i in range(num_cases)
|
||||
]
|
||||
|
||||
oracle = HashOracle(vocab_size=vocab_size)
|
||||
out = oracle.expected_tokens(
|
||||
generalized_req_ids=torch.tensor(generalized_req_ids, dtype=torch.int64),
|
||||
positions=torch.tensor(positions, dtype=torch.int64),
|
||||
)
|
||||
torch_tokens = out.tolist()
|
||||
|
||||
for i in range(num_cases):
|
||||
self.assertEqual(
|
||||
torch_tokens[i],
|
||||
ref_tokens[i],
|
||||
f"batched mismatch at case {i}: generalized_req_id={generalized_req_ids[i]} "
|
||||
f"position={positions[i]}: torch={torch_tokens[i]} ref={ref_tokens[i]}",
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,49 @@
|
||||
"""install_oracle_sampler registration into sglang's sampler-backend registry.
|
||||
|
||||
Instantiating the registered _OracleSampler factory requires a live distributed (TP) group
|
||||
plus a populated global ServerArgs, so the forward-path behavior of _OracleSampler is covered
|
||||
by the e2e harness rather than this unit file. Here we only assert the registration-side
|
||||
contract: the backend name shows up in the registry / choice set, and second install replaces
|
||||
the factory with one bound to the new oracle.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import unittest
|
||||
|
||||
os.environ["SGLANG_KV_CANARY_ENABLE_TOKEN_ORACLE"] = "1"
|
||||
|
||||
from sglang.srt.kv_canary.token_oracle.oracle import HashOracle
|
||||
from sglang.srt.kv_canary.token_oracle.sampler import install_oracle_sampler
|
||||
from sglang.srt.layers.sampler import _CUSTOM_SAMPLER_FACTORIES
|
||||
from sglang.srt.server_args import SAMPLING_BACKEND_CHOICES
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cuda_ci(est_time=60, stage="extra-a", runner_config="1-gpu-small")
|
||||
|
||||
|
||||
class TestInstallOracleSampler(CustomTestCase):
|
||||
def test_install_oracle_sampler_twice_returns_distinct_hooks_with_replaced_oracle(
|
||||
self,
|
||||
) -> None:
|
||||
"""Verify reinstalling the oracle sampler replaces the registered factory."""
|
||||
oracle_a = HashOracle(vocab_size=100)
|
||||
oracle_b = HashOracle(vocab_size=100)
|
||||
|
||||
hook_a = install_oracle_sampler(oracle=oracle_a)
|
||||
self.assertIn("token_oracle", _CUSTOM_SAMPLER_FACTORIES)
|
||||
self.assertIn("token_oracle", SAMPLING_BACKEND_CHOICES)
|
||||
factory_a = _CUSTOM_SAMPLER_FACTORIES["token_oracle"]
|
||||
self.assertIs(hook_a.oracle, oracle_a)
|
||||
|
||||
hook_b = install_oracle_sampler(oracle=oracle_b)
|
||||
factory_b = _CUSTOM_SAMPLER_FACTORIES["token_oracle"]
|
||||
self.assertIs(hook_b.oracle, oracle_b)
|
||||
self.assertIsNot(hook_a, hook_b)
|
||||
self.assertIsNot(factory_a, factory_b)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,8 +1,11 @@
|
||||
import importlib
|
||||
import json
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import sglang.srt.server_args as server_args_module
|
||||
from sglang.srt.arg_groups.speculative_hook import handle_speculative_decoding
|
||||
from sglang.srt.server_args import PortArgs, ServerArgs, prepare_server_args
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
@@ -659,5 +662,58 @@ class TestCutedslMoeMaxNumTokens(unittest.TestCase):
|
||||
self.assertEqual(args.cutedsl_moe_max_num_tokens(), 512)
|
||||
|
||||
|
||||
class TestSamplingBackendTokenOracleEnvGate(CustomTestCase):
|
||||
"""The 'token_oracle' choice is gated on SGLANG_KV_CANARY_ENABLE_TOKEN_ORACLE.
|
||||
|
||||
The choice set is built once at server_args.py import time, so each subtest
|
||||
reloads the module with the env var set to the desired value.
|
||||
"""
|
||||
|
||||
def _reload_server_args_with_env(self, *, enabled: bool):
|
||||
previous = os.environ.get("SGLANG_KV_CANARY_ENABLE_TOKEN_ORACLE")
|
||||
os.environ["SGLANG_KV_CANARY_ENABLE_TOKEN_ORACLE"] = "1" if enabled else "0"
|
||||
try:
|
||||
return importlib.reload(server_args_module)
|
||||
finally:
|
||||
if previous is None:
|
||||
os.environ.pop("SGLANG_KV_CANARY_ENABLE_TOKEN_ORACLE", None)
|
||||
else:
|
||||
os.environ["SGLANG_KV_CANARY_ENABLE_TOKEN_ORACLE"] = previous
|
||||
|
||||
def test_token_oracle_rejected_when_env_disabled(self):
|
||||
reloaded = self._reload_server_args_with_env(enabled=False)
|
||||
self.assertNotIn("token_oracle", reloaded.SAMPLING_BACKEND_CHOICES)
|
||||
|
||||
with self.assertRaises(SystemExit):
|
||||
reloaded.prepare_server_args(
|
||||
[
|
||||
"--model-path",
|
||||
DEFAULT_SMALL_MODEL_NAME_FOR_TEST_QWEN,
|
||||
"--sampling-backend",
|
||||
"token_oracle",
|
||||
]
|
||||
)
|
||||
|
||||
def test_token_oracle_accepted_when_env_enabled(self):
|
||||
reloaded = self._reload_server_args_with_env(enabled=True)
|
||||
self.assertIn("token_oracle", reloaded.SAMPLING_BACKEND_CHOICES)
|
||||
|
||||
parsed = reloaded.prepare_server_args(
|
||||
[
|
||||
"--model-path",
|
||||
DEFAULT_SMALL_MODEL_NAME_FOR_TEST_QWEN,
|
||||
"--sampling-backend",
|
||||
"token_oracle",
|
||||
# Explicit device so ServerArgs.__post_init__ does not call
|
||||
# get_device() (fails on CPU-only CI runners) and does not run
|
||||
# _handle_cpu_backends (which would override sampling_backend
|
||||
# to "pytorch", masking what we want to verify).
|
||||
"--device",
|
||||
"cuda",
|
||||
]
|
||||
)
|
||||
self.assertEqual(parsed.sampling_backend, "token_oracle")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
Reference in New Issue
Block a user