Add a deterministic token oracle and production write-input assertion (#26815)
This commit is contained in:
fzyzcjy
@@ -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()
|
||||
Reference in New Issue
Block a user