7ee8a6dd01
Signed-off-by: Raphael Rialland <raphael.rialland@mistral.ai> Signed-off-by: Simon Veitner <sveitner@redhat.com> Signed-off-by: Tomas Ruiz <tomas.ruiz.te@gmail.com> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: Simon Veitner <sveitner@redhat.com> Co-authored-by: Tomas Ruiz <tomas.ruiz.te@gmail.com>
1029 lines
35 KiB
Python
1029 lines
35 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from types import SimpleNamespace
|
|
|
|
import numpy as np
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm import SamplingParams
|
|
from vllm.config.watermarking import WatermarkConfig
|
|
from vllm.platforms import current_platform
|
|
from vllm.v1.watermarking import (
|
|
DualKeyGumbelWatermarker,
|
|
SupportsSpeculativeDecoding,
|
|
create_watermarker,
|
|
derive_watermark_key,
|
|
)
|
|
from vllm.v1.watermarking.gpu_sampler import GPUWatermarkSampler
|
|
from vllm.v1.watermarking.gumbel import GumbelWatermarker
|
|
from vllm.v1.watermarking.spec_decode import (
|
|
DraftWatermarker,
|
|
_resolve_watermark_key,
|
|
create_speculative_draft_watermarker,
|
|
create_speculative_target_watermarker,
|
|
speculative_target_watermark_key,
|
|
)
|
|
from vllm.v1.watermarking.watermarker import Watermarker, WatermarkSample
|
|
from vllm.v1.worker.gpu.sample.sampler import Sampler
|
|
from vllm.v1.worker.gpu.sample.watermark import (
|
|
philox_gumbel_sample,
|
|
repeated_context_mask,
|
|
)
|
|
from vllm.v1.worker.gpu.spec_decode.dspark.speculator import DSparkSpeculator
|
|
from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator
|
|
|
|
|
|
class StubWatermarker(Watermarker):
|
|
context_width = 1
|
|
|
|
def _sample_watermarked(self, logits, contexts):
|
|
return WatermarkSample(torch.tensor([7, 7]), logits + 10)
|
|
|
|
|
|
def _argmax_sampler(logits: torch.Tensor) -> torch.Tensor:
|
|
return logits.argmax(dim=-1)
|
|
|
|
|
|
@pytest.mark.parametrize("algorithm", ["gumbel", "dual_key_gumbel"])
|
|
def test_watermarker_contract(algorithm: str):
|
|
watermarker = create_watermarker(
|
|
WatermarkConfig(algorithm=algorithm, key=42, context_width=4)
|
|
)
|
|
logits = torch.zeros(2, 128)
|
|
contexts = torch.tensor([[1, 2, 3, 4], [4, 5, 6, 7]])
|
|
|
|
first = watermarker.sample(logits, contexts, _argmax_sampler)
|
|
second = watermarker.sample(logits, contexts, _argmax_sampler)
|
|
|
|
assert first.token_ids.shape == (2,)
|
|
assert first.logits.shape == logits.shape
|
|
assert torch.equal(first.token_ids, second.token_ids)
|
|
|
|
|
|
@pytest.mark.parametrize("algorithm", ["gumbel", "dual_key_gumbel"])
|
|
@pytest.mark.parametrize(
|
|
"config_overrides",
|
|
[
|
|
{"deduplicate_contexts": "none"},
|
|
{"deduplicate_contexts_max_history": 1023},
|
|
],
|
|
)
|
|
def test_gumbel_config_warns_when_context_deduplication_is_weak(
|
|
monkeypatch, algorithm, config_overrides
|
|
):
|
|
messages: list[str] = []
|
|
monkeypatch.setattr(
|
|
"vllm.config.watermarking.logger.warning_once",
|
|
lambda message, *, scope: messages.append(message),
|
|
)
|
|
|
|
WatermarkConfig(key=42, algorithm=algorithm)
|
|
WatermarkConfig(key=42, algorithm=algorithm, deduplicate_contexts_max_history=1024)
|
|
WatermarkConfig(key=42, algorithm=algorithm, deduplicate_contexts_max_history=None)
|
|
WatermarkConfig(key=42, algorithm=algorithm, **config_overrides)
|
|
|
|
assert messages == [
|
|
(
|
|
"Gumbel-max watermarking with context deduplication disabled "
|
|
"or limited to fewer than 1024 positions may increase the frequency of "
|
|
"degenerate generations, including repetition loops. Use "
|
|
"deduplicate_contexts='single_turn' or 'all' with "
|
|
"deduplicate_contexts_max_history at least 1024 or null to mitigate this."
|
|
)
|
|
]
|
|
|
|
|
|
def test_context_deduplication_history_can_be_unbounded():
|
|
config = WatermarkConfig(key=42, deduplicate_contexts_max_history=None)
|
|
|
|
assert config.deduplicate_contexts_max_history is None
|
|
|
|
|
|
def test_large_context_width_warns_but_is_allowed():
|
|
config = WatermarkConfig(key=42, context_width=17)
|
|
|
|
with pytest.warns(UserWarning, match="reduce robustness to edits"):
|
|
watermarker = create_watermarker(config)
|
|
|
|
assert watermarker.context_width == 17
|
|
|
|
|
|
def test_dual_key_watermarker_uses_domain_separated_keys():
|
|
config = WatermarkConfig(algorithm="dual_key_gumbel", key=42)
|
|
|
|
watermarker = create_watermarker(config)
|
|
assert isinstance(watermarker, SupportsSpeculativeDecoding)
|
|
draft = watermarker.draft_watermarker
|
|
target = watermarker.target_watermarker
|
|
|
|
assert isinstance(watermarker, DualKeyGumbelWatermarker)
|
|
assert watermarker.alpha == 0.1
|
|
assert draft.prf.key == derive_watermark_key(42, b"key_a")
|
|
assert target.prf.key == derive_watermark_key(42, b"key_b")
|
|
assert target.prf.key != draft.prf.key
|
|
|
|
|
|
def test_dual_key_watermarker_routes_tokens_with_alpha():
|
|
watermarker = DualKeyGumbelWatermarker(key=42, context_width=2, alpha=0.25)
|
|
logits = torch.zeros(2, 16)
|
|
contexts = torch.tensor([[1, 2], [3, 4]])
|
|
key_a = watermarker.draft_watermarker.sample(logits, contexts, lambda _: None)
|
|
key_b = watermarker.target_watermarker.sample(logits, contexts, lambda _: None)
|
|
|
|
def route(routing_logits):
|
|
torch.testing.assert_close(
|
|
routing_logits.softmax(dim=-1),
|
|
torch.tensor([[0.75, 0.25], [0.75, 0.25]]),
|
|
)
|
|
return torch.tensor([0, 1])
|
|
|
|
sampled = watermarker.sample(logits, contexts, route)
|
|
|
|
assert torch.equal(
|
|
sampled.token_ids, torch.stack([key_a.token_ids[0], key_b.token_ids[1]])
|
|
)
|
|
|
|
|
|
def test_dual_key_routing_logits_are_cached():
|
|
alpha = 0.25
|
|
watermarker = DualKeyGumbelWatermarker(key=42, context_width=2, alpha=alpha)
|
|
logits = torch.zeros(2, 16)
|
|
contexts = torch.tensor([[1, 2], [3, 4]])
|
|
recorded: list[torch.Tensor] = []
|
|
|
|
def route(routing_logits):
|
|
recorded.append(routing_logits)
|
|
return torch.tensor([0, 1])
|
|
|
|
watermarker.sample(logits, contexts, route)
|
|
watermarker.sample(logits, contexts, route)
|
|
|
|
expected = torch.tensor([1 - alpha, alpha], dtype=torch.float32).log()
|
|
expected = expected.expand(2, -1)
|
|
assert len(recorded) == 2
|
|
assert all(torch.equal(seen, expected) for seen in recorded)
|
|
assert len(watermarker._routing_logits_cache) == 1
|
|
|
|
|
|
@pytest.mark.parametrize("alpha", [0.0, 0.25, 1.0])
|
|
def test_dual_key_watermarker_skips_masked_rows(alpha: float):
|
|
watermarker = DualKeyGumbelWatermarker(key=42, context_width=2, alpha=alpha)
|
|
torch.manual_seed(0)
|
|
logits = torch.randn(4, 32)
|
|
contexts = torch.tensor([[1, 2], [3, 4], [5, 6], [7, 8]])
|
|
skip_mask = torch.tensor([True, False, True, False])
|
|
|
|
unmasked = watermarker.sample(logits, contexts, _argmax_sampler)
|
|
masked = watermarker.sample(logits, contexts, _argmax_sampler, skip_mask)
|
|
|
|
ordinary = _argmax_sampler(logits)
|
|
assert torch.equal(masked.token_ids[skip_mask], ordinary[skip_mask])
|
|
assert torch.equal(masked.token_ids[~skip_mask], unmasked.token_ids[~skip_mask])
|
|
assert not torch.equal(unmasked.token_ids[skip_mask], ordinary[skip_mask])
|
|
|
|
|
|
def test_dual_key_watermarker_routing_requires_a_random_sampler():
|
|
watermarker = DualKeyGumbelWatermarker(key=42, context_width=2, alpha=0.25)
|
|
|
|
with pytest.raises(ValueError, match="random sampler"):
|
|
watermarker.sample(torch.zeros(2, 16), torch.tensor([[1, 2], [3, 4]]))
|
|
|
|
|
|
def test_speculative_decoding_uses_fixed_dual_key_roles():
|
|
watermarker = create_watermarker(
|
|
WatermarkConfig(algorithm="dual_key_gumbel", key=42, alpha=0.25)
|
|
)
|
|
|
|
target = create_speculative_target_watermarker(watermarker)
|
|
draft = create_speculative_draft_watermarker(
|
|
watermarker,
|
|
max_num_reqs=1,
|
|
device=torch.device("cpu"),
|
|
allow_target_only=False,
|
|
)
|
|
|
|
assert target.prf.key == derive_watermark_key(42, b"key_b")
|
|
assert draft is not None
|
|
assert draft.watermarker.prf.key == derive_watermark_key(42, b"key_a")
|
|
|
|
|
|
def test_recovery_key_rejects_the_unsplit_dual_key_watermarker():
|
|
watermarker = create_watermarker(
|
|
WatermarkConfig(algorithm="dual_key_gumbel", key=42, alpha=0.25)
|
|
)
|
|
|
|
with pytest.raises(ValueError, match="keys the target role separately"):
|
|
_resolve_watermark_key(watermarker)
|
|
|
|
target = create_speculative_target_watermarker(watermarker)
|
|
assert _resolve_watermark_key(target) == derive_watermark_key(42, b"key_b")
|
|
|
|
|
|
@pytest.mark.parametrize("algorithm", ["gumbel", "dual_key_gumbel"])
|
|
def test_config_key_resolution_matches_the_model_runner(algorithm: str):
|
|
config = WatermarkConfig(algorithm=algorithm, key=42, alpha=0.25)
|
|
runtime_key = _resolve_watermark_key(
|
|
create_speculative_target_watermarker(create_watermarker(config))
|
|
)
|
|
|
|
assert speculative_target_watermark_key(config) == runtime_key
|
|
|
|
|
|
def test_config_key_resolution_returns_none_when_watermarking_is_disabled():
|
|
assert speculative_target_watermark_key(None) is None
|
|
|
|
|
|
def test_target_only_speculative_watermarking_skips_draft_watermarker():
|
|
watermarker = create_watermarker(WatermarkConfig(algorithm="gumbel", key=42))
|
|
|
|
assert not isinstance(watermarker, SupportsSpeculativeDecoding)
|
|
with pytest.raises(ValueError, match="does not support speculative decoding"):
|
|
create_speculative_draft_watermarker(
|
|
watermarker,
|
|
max_num_reqs=1,
|
|
device=torch.device("cpu"),
|
|
allow_target_only=False,
|
|
)
|
|
assert (
|
|
create_speculative_draft_watermarker(
|
|
watermarker,
|
|
max_num_reqs=1,
|
|
device=torch.device("cpu"),
|
|
allow_target_only=True,
|
|
)
|
|
is None
|
|
)
|
|
|
|
|
|
def test_dual_key_derivation_is_stable():
|
|
assert derive_watermark_key(32, b"key_a") == 16368605726115524094
|
|
assert derive_watermark_key(32, b"key_b") == 4799302812959726346
|
|
|
|
|
|
def test_sampling_params_can_disable_watermarking():
|
|
assert SamplingParams().watermarking
|
|
assert not SamplingParams.from_optional(watermarking=False).watermarking
|
|
|
|
|
|
def test_gpu_sampler_warns_when_watermarking_is_enabled_for_greedy(monkeypatch):
|
|
sampler = object.__new__(GPUWatermarkSampler)
|
|
sampler.watermarking = SimpleNamespace(np=np.ones(1, dtype=bool))
|
|
messages: list[str] = []
|
|
monkeypatch.setattr(Sampler, "add_request", lambda *args, **kwargs: None)
|
|
monkeypatch.setattr(
|
|
"vllm.v1.watermarking.gpu_sampler.logger.warning_once", messages.append
|
|
)
|
|
|
|
sampler.add_request(0, 1, SamplingParams(temperature=0))
|
|
sampler.add_request(0, 1, SamplingParams(temperature=1))
|
|
sampler.add_request(0, 1, SamplingParams(temperature=0, watermarking=False))
|
|
|
|
assert messages == [
|
|
(
|
|
"Watermarking is enabled, but greedy decoding (temperature=0) cannot be "
|
|
"watermarked. This request will use ordinary greedy sampling."
|
|
)
|
|
]
|
|
|
|
|
|
def test_gpu_sampler_respects_mixed_request_watermarking(monkeypatch):
|
|
sampler = object.__new__(GPUWatermarkSampler)
|
|
sampler.watermarker = StubWatermarker()
|
|
sampler.deduplicate_contexts = "single_turn"
|
|
sampler.watermarking = SimpleNamespace(
|
|
np=np.array([True, False]), gpu=torch.tensor([True, False])
|
|
)
|
|
sampler.sampling_states = SimpleNamespace(
|
|
temperature=SimpleNamespace(np=np.ones(2), gpu=torch.ones(2)),
|
|
seeds=SimpleNamespace(gpu=torch.zeros(2, dtype=torch.int64)),
|
|
)
|
|
sampler.use_fp64_gumbel = False
|
|
sampler._get_contexts = lambda expanded_idx_mapping: torch.zeros(
|
|
2, 1, dtype=torch.int64
|
|
)
|
|
sampler._get_repeated_contexts = lambda expanded_idx_mapping, contexts: torch.zeros(
|
|
2, dtype=torch.bool
|
|
)
|
|
monkeypatch.setattr(
|
|
"vllm.v1.watermarking.watermarker.gumbel_sample",
|
|
lambda *args, **kwargs: torch.tensor([3, 4]),
|
|
)
|
|
logits = torch.zeros(2, 8)
|
|
|
|
sampled, output_logits = sampler._sample_random(
|
|
logits,
|
|
torch.tensor([0, 1]),
|
|
np.array([0, 1]),
|
|
torch.zeros(2, dtype=torch.int64),
|
|
None,
|
|
None,
|
|
False,
|
|
)
|
|
|
|
assert torch.equal(sampled, torch.tensor([7, 4]))
|
|
assert torch.equal(output_logits[0], torch.full((8,), 10.0))
|
|
assert torch.equal(output_logits[1], logits[1])
|
|
|
|
|
|
def test_gpu_sampler_skips_watermarking_for_repeated_contexts(monkeypatch):
|
|
sampler = object.__new__(GPUWatermarkSampler)
|
|
sampler.watermarker = StubWatermarker()
|
|
sampler.deduplicate_contexts = "single_turn"
|
|
sampler.watermarking = SimpleNamespace(
|
|
np=np.array([True, True]), gpu=torch.tensor([True, True])
|
|
)
|
|
sampler.sampling_states = SimpleNamespace(
|
|
temperature=SimpleNamespace(np=np.ones(2), gpu=torch.ones(2)),
|
|
seeds=SimpleNamespace(gpu=torch.zeros(2, dtype=torch.int64)),
|
|
)
|
|
sampler.use_fp64_gumbel = False
|
|
sampler._get_contexts = lambda expanded_idx_mapping: torch.zeros(
|
|
2, 1, dtype=torch.int64
|
|
)
|
|
sampler._get_repeated_contexts = lambda expanded_idx_mapping, contexts: (
|
|
torch.tensor([True, False])
|
|
)
|
|
monkeypatch.setattr(
|
|
"vllm.v1.watermarking.watermarker.gumbel_sample",
|
|
lambda *args, **kwargs: torch.tensor([3, 4]),
|
|
)
|
|
logits = torch.zeros(2, 8)
|
|
|
|
sampled, output_logits = sampler._sample_random(
|
|
logits,
|
|
torch.tensor([0, 1]),
|
|
np.array([0, 1]),
|
|
torch.zeros(2, dtype=torch.int64),
|
|
None,
|
|
None,
|
|
False,
|
|
)
|
|
|
|
assert torch.equal(sampled, torch.tensor([3, 7]))
|
|
assert torch.equal(output_logits[0], logits[0])
|
|
assert torch.equal(output_logits[1], torch.full((8,), 10.0))
|
|
|
|
|
|
def test_gpu_sampler_can_disable_context_deduplication(monkeypatch):
|
|
sampler = object.__new__(GPUWatermarkSampler)
|
|
sampler.watermarker = StubWatermarker()
|
|
sampler.deduplicate_contexts = "none"
|
|
sampler.watermarking = SimpleNamespace(
|
|
np=np.array([True, True]), gpu=torch.tensor([True, True])
|
|
)
|
|
sampler.sampling_states = SimpleNamespace(
|
|
temperature=SimpleNamespace(np=np.ones(2), gpu=torch.ones(2)),
|
|
seeds=SimpleNamespace(gpu=torch.zeros(2, dtype=torch.int64)),
|
|
)
|
|
sampler.use_fp64_gumbel = False
|
|
sampler._get_contexts = lambda expanded_idx_mapping: torch.zeros(
|
|
2, 1, dtype=torch.int64
|
|
)
|
|
sampler._get_repeated_contexts = lambda *args: pytest.fail(
|
|
"context deduplication should not run"
|
|
)
|
|
logits = torch.zeros(2, 8)
|
|
|
|
sampled, output_logits = sampler._sample_random(
|
|
logits,
|
|
torch.tensor([0, 1]),
|
|
np.array([0, 1]),
|
|
torch.zeros(2, dtype=torch.int64),
|
|
None,
|
|
None,
|
|
False,
|
|
)
|
|
|
|
assert torch.equal(sampled, torch.tensor([7, 7]))
|
|
assert torch.equal(output_logits, torch.full((2, 8), 10.0))
|
|
|
|
|
|
def test_repeated_context_mask_ignores_prompt_tokens():
|
|
all_token_ids = torch.tensor(
|
|
[
|
|
[8, 9, 1, 2, 1, 2],
|
|
[3, 4, 1, 2, 3, 4],
|
|
[0, 0, 0, 0, 0, 0],
|
|
],
|
|
dtype=torch.int32,
|
|
)
|
|
req_indices = torch.tensor([0, 1, -1])
|
|
prompt_lens = torch.tensor([2, 2, 0])
|
|
total_lens = torch.tensor([6, 6, 0])
|
|
contexts = torch.tensor([[1, 2], [3, 4], [-1, -1]])
|
|
|
|
repeated = repeated_context_mask(
|
|
all_token_ids,
|
|
req_indices,
|
|
prompt_lens,
|
|
total_lens,
|
|
contexts,
|
|
)
|
|
|
|
assert torch.equal(repeated, torch.tensor([True, False, False]))
|
|
|
|
|
|
def test_repeated_context_mask_can_include_prompt_tokens():
|
|
all_token_ids = torch.tensor([[1, 2, 3, 4, 5, 6, 1, 2, 3, 4]])
|
|
req_indices = torch.tensor([0])
|
|
prompt_lens = torch.tensor([4])
|
|
total_lens = torch.tensor([10])
|
|
contexts = torch.tensor([[1, 2, 3, 4]])
|
|
|
|
single_turn = repeated_context_mask(
|
|
all_token_ids, req_indices, prompt_lens, total_lens, contexts
|
|
)
|
|
all_history = repeated_context_mask(
|
|
all_token_ids,
|
|
req_indices,
|
|
prompt_lens,
|
|
total_lens,
|
|
contexts,
|
|
include_prompt=True,
|
|
)
|
|
|
|
assert not single_turn.item()
|
|
assert all_history.item()
|
|
|
|
|
|
def test_repeated_context_mask_can_skip_partial_contexts():
|
|
all_token_ids = torch.tensor([[10, 11, 12, 13, 1, 2, 3, 4]], dtype=torch.int32)
|
|
req_indices = torch.tensor([0])
|
|
prompt_lens = torch.tensor([4])
|
|
contexts = torch.tensor([[-1, -1, -1, 1]])
|
|
|
|
partial = repeated_context_mask(
|
|
all_token_ids,
|
|
req_indices,
|
|
prompt_lens,
|
|
torch.tensor([5]),
|
|
contexts,
|
|
include_prompt=True,
|
|
skip_partial_context=True,
|
|
)
|
|
complete = repeated_context_mask(
|
|
all_token_ids,
|
|
req_indices,
|
|
prompt_lens,
|
|
torch.tensor([8]),
|
|
torch.tensor([[1, 2, 3, 4]]),
|
|
include_prompt=True,
|
|
skip_partial_context=True,
|
|
)
|
|
|
|
assert partial.item()
|
|
assert not complete.item()
|
|
|
|
|
|
def test_gpu_sampler_contexts_are_completion_local_for_all_scopes():
|
|
sampler = object.__new__(GPUWatermarkSampler)
|
|
sampler.watermarker = SimpleNamespace(context_width=4)
|
|
sampler.req_states = SimpleNamespace(
|
|
all_token_ids=SimpleNamespace(gpu=torch.tensor([[10, 11, 12, 13, 1, 2]])),
|
|
prompt_len=SimpleNamespace(gpu=torch.tensor([4])),
|
|
total_len=SimpleNamespace(gpu=torch.tensor([6])),
|
|
)
|
|
request_indices = torch.tensor([0])
|
|
|
|
sampler.deduplicate_contexts = "single_turn"
|
|
single_turn = sampler._get_contexts(request_indices)
|
|
sampler.deduplicate_contexts = "all"
|
|
all_history = sampler._get_contexts(request_indices)
|
|
|
|
assert torch.equal(single_turn, torch.tensor([[-1, -1, 1, 2]]))
|
|
assert torch.equal(all_history, single_turn)
|
|
|
|
|
|
def test_gpu_sampler_all_history_skips_partial_context():
|
|
sampler = object.__new__(GPUWatermarkSampler)
|
|
sampler.watermarker = SimpleNamespace(context_width=4)
|
|
sampler.req_states = SimpleNamespace(
|
|
all_token_ids=SimpleNamespace(gpu=torch.tensor([[10, 11, 12, 13, 1, 2]])),
|
|
prompt_len=SimpleNamespace(gpu=torch.tensor([4])),
|
|
total_len=SimpleNamespace(gpu=torch.tensor([6])),
|
|
)
|
|
sampler.deduplicate_contexts_max_history = None
|
|
request_indices = torch.tensor([0])
|
|
contexts = sampler._get_contexts(request_indices)
|
|
|
|
sampler.deduplicate_contexts = "single_turn"
|
|
single_turn = sampler._get_repeated_contexts(request_indices, contexts)
|
|
sampler.deduplicate_contexts = "all"
|
|
all_history = sampler._get_repeated_contexts(request_indices, contexts)
|
|
|
|
assert not single_turn.item()
|
|
assert all_history.item()
|
|
|
|
|
|
def test_repeated_context_mask_respects_max_history():
|
|
all_token_ids = torch.tensor([[1, 2, 3, 4, 5, 6, 1, 2, 3, 4]])
|
|
req_indices = torch.tensor([0])
|
|
prompt_lens = torch.tensor([0])
|
|
total_lens = torch.tensor([10])
|
|
contexts = torch.tensor([[1, 2, 3, 4]])
|
|
|
|
full_history = repeated_context_mask(
|
|
all_token_ids, req_indices, prompt_lens, total_lens, contexts
|
|
)
|
|
last_six = repeated_context_mask(
|
|
all_token_ids, req_indices, prompt_lens, total_lens, contexts, max_history=6
|
|
)
|
|
last_five = repeated_context_mask(
|
|
all_token_ids, req_indices, prompt_lens, total_lens, contexts, max_history=5
|
|
)
|
|
|
|
assert full_history.item()
|
|
assert last_six.item()
|
|
assert not last_five.item()
|
|
|
|
|
|
@pytest.mark.parametrize("max_history", [0, -1])
|
|
def test_repeated_context_mask_rejects_non_positive_max_history(max_history: int):
|
|
with pytest.raises(ValueError, match="max_history must be positive or None"):
|
|
repeated_context_mask(
|
|
torch.tensor([[1, 2, 3, 4]], dtype=torch.int32),
|
|
torch.tensor([0], dtype=torch.int32),
|
|
torch.tensor([0], dtype=torch.int32),
|
|
torch.tensor([4], dtype=torch.int32),
|
|
torch.tensor([[1, 2, 3, 4]], dtype=torch.int32),
|
|
max_history=max_history,
|
|
)
|
|
|
|
|
|
def test_gpu_sampler_wires_request_history_scope():
|
|
sampler = object.__new__(GPUWatermarkSampler)
|
|
sampler.req_states = SimpleNamespace(
|
|
all_token_ids=SimpleNamespace(
|
|
gpu=torch.tensor([[1, 2, 3, 4, 5, 6, 1, 2, 3, 4]])
|
|
),
|
|
prompt_len=SimpleNamespace(gpu=torch.tensor([4])),
|
|
total_len=SimpleNamespace(gpu=torch.tensor([10])),
|
|
)
|
|
sampler.deduplicate_contexts_max_history = None
|
|
request_indices = torch.tensor([0])
|
|
contexts = torch.tensor([[1, 2, 3, 4]])
|
|
|
|
sampler.deduplicate_contexts = "single_turn"
|
|
single_turn = sampler._get_repeated_contexts(request_indices, contexts)
|
|
sampler.deduplicate_contexts = "all"
|
|
all_history = sampler._get_repeated_contexts(request_indices, contexts)
|
|
|
|
assert not single_turn.item()
|
|
assert all_history.item()
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator"
|
|
)
|
|
@pytest.mark.parametrize("max_history", [5, 6])
|
|
@pytest.mark.parametrize("include_prompt", [False, True])
|
|
def test_repeated_context_mask_max_history_accelerator_parity(
|
|
max_history: int, include_prompt: bool
|
|
):
|
|
inputs = (
|
|
torch.tensor([[1, 2, 3, 4, 5, 6, 1, 2, 3, 4]]),
|
|
torch.tensor([0]),
|
|
torch.tensor([0]),
|
|
torch.tensor([10]),
|
|
torch.tensor([[1, 2, 3, 4]]),
|
|
)
|
|
|
|
expected = repeated_context_mask(
|
|
*inputs, max_history=max_history, include_prompt=include_prompt
|
|
)
|
|
actual = repeated_context_mask(
|
|
*(value.cuda() for value in inputs),
|
|
max_history=max_history,
|
|
include_prompt=include_prompt,
|
|
).cpu()
|
|
|
|
assert torch.equal(actual, expected)
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator"
|
|
)
|
|
def test_repeated_context_mask_partial_context_accelerator_parity():
|
|
inputs = (
|
|
torch.tensor([[10, 11, 12, 13, 1]], dtype=torch.int32),
|
|
torch.tensor([0], dtype=torch.int32),
|
|
torch.tensor([4], dtype=torch.int32),
|
|
torch.tensor([5], dtype=torch.int32),
|
|
torch.tensor([[-1, -1, -1, 1]], dtype=torch.int64),
|
|
)
|
|
|
|
expected = repeated_context_mask(
|
|
*inputs, include_prompt=True, skip_partial_context=True
|
|
)
|
|
actual = repeated_context_mask(
|
|
*(value.cuda() for value in inputs),
|
|
include_prompt=True,
|
|
skip_partial_context=True,
|
|
).cpu()
|
|
|
|
assert torch.equal(actual, expected)
|
|
|
|
|
|
def _unaligned_max_history_inputs() -> tuple[torch.Tensor, ...]:
|
|
"""One row whose only repeated context ends at position 680.
|
|
|
|
680 is not a multiple of the kernel's 512-position block, so the scan window
|
|
selected by ``max_history`` starts in the middle of a block.
|
|
"""
|
|
history = list(range(1_000, 2_200))
|
|
history[676:680] = [1, 2, 3, 4]
|
|
return (
|
|
torch.tensor([history], dtype=torch.int32),
|
|
torch.tensor([0], dtype=torch.int32),
|
|
torch.tensor([0], dtype=torch.int32),
|
|
torch.tensor([len(history)], dtype=torch.int32),
|
|
torch.tensor([[1, 2, 3, 4]], dtype=torch.int32),
|
|
)
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator"
|
|
)
|
|
@pytest.mark.parametrize(
|
|
"max_history,expected", [(1200 - 680, True), (1200 - 680 - 1, False)]
|
|
)
|
|
def test_repeated_context_mask_unaligned_max_history_accelerator_parity(
|
|
max_history: int, expected: bool
|
|
):
|
|
inputs = _unaligned_max_history_inputs()
|
|
|
|
reference = repeated_context_mask(*inputs, max_history=max_history)
|
|
actual = repeated_context_mask(
|
|
*(value.cuda() for value in inputs), max_history=max_history
|
|
).cpu()
|
|
|
|
assert torch.equal(reference, torch.tensor([expected]))
|
|
assert torch.equal(actual, reference)
|
|
|
|
|
|
def _repeated_context_inputs(
|
|
context_width: int, device: str = "cpu"
|
|
) -> tuple[torch.Tensor, ...]:
|
|
"""Build repeated, unique, and padded request rows for mask tests."""
|
|
prompt = [101, 102, 103, 104]
|
|
repeated_output = [*range(1, context_width + 1)] * 2
|
|
unique_output = [*range(1, context_width + 2)]
|
|
max_len = len(prompt) + len(repeated_output)
|
|
rows = [
|
|
prompt + repeated_output,
|
|
prompt + unique_output,
|
|
[],
|
|
]
|
|
all_token_ids = torch.zeros((3, max_len), dtype=torch.int32, device=device)
|
|
for row, token_ids in enumerate(rows):
|
|
all_token_ids[row, : len(token_ids)] = torch.tensor(
|
|
token_ids, dtype=torch.int32, device=device
|
|
)
|
|
return (
|
|
all_token_ids,
|
|
torch.tensor([0, 1, -1], dtype=torch.int32, device=device),
|
|
torch.tensor([len(prompt), len(prompt), 0], dtype=torch.int32, device=device),
|
|
torch.tensor([len(rows[0]), len(rows[1]), 0], dtype=torch.int32, device=device),
|
|
torch.tensor(
|
|
[
|
|
repeated_output[-context_width:],
|
|
unique_output[-context_width:],
|
|
[-1] * context_width,
|
|
],
|
|
dtype=torch.int64,
|
|
device=device,
|
|
),
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("context_width", [1, 3, 4, 16, 17])
|
|
def test_repeated_context_mask(context_width: int):
|
|
repeated = repeated_context_mask(*_repeated_context_inputs(context_width))
|
|
|
|
assert torch.equal(repeated, torch.tensor([True, False, False]))
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator"
|
|
)
|
|
@pytest.mark.parametrize("context_width", [1, 3, 4, 16, 17])
|
|
def test_repeated_context_mask_accelerator_parity(context_width: int):
|
|
cpu_inputs = _repeated_context_inputs(context_width)
|
|
accelerator_inputs = list(_repeated_context_inputs(context_width, "cuda"))
|
|
contexts = accelerator_inputs[-1]
|
|
storage = torch.empty(
|
|
contexts.shape[0], contexts.shape[1] * 2, dtype=contexts.dtype, device="cuda"
|
|
)
|
|
storage[:, ::2] = contexts
|
|
accelerator_inputs[-1] = storage[:, ::2]
|
|
|
|
expected = repeated_context_mask(*cpu_inputs)
|
|
actual = repeated_context_mask(*accelerator_inputs).cpu()
|
|
|
|
assert torch.equal(actual, expected)
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator"
|
|
)
|
|
def test_repeated_context_mask_scans_multiple_blocks():
|
|
context = [1, 2, 3, 4]
|
|
repeated_output = [*range(10_000, 11_030), *context, 99, *context]
|
|
unique_output = list(range(20_000, 21_039))
|
|
prompt = [101, 102, 103, 104]
|
|
rows = [prompt + repeated_output, prompt + unique_output]
|
|
all_token_ids = torch.zeros((2, len(rows[0])), dtype=torch.int32)
|
|
for row, token_ids in enumerate(rows):
|
|
all_token_ids[row, : len(token_ids)] = torch.tensor(token_ids)
|
|
inputs = (
|
|
all_token_ids,
|
|
torch.tensor([0, 1], dtype=torch.int32),
|
|
torch.tensor([len(prompt), len(prompt)], dtype=torch.int32),
|
|
torch.tensor([len(rows[0]), len(rows[1])], dtype=torch.int32),
|
|
torch.tensor([context, unique_output[-4:]], dtype=torch.int64),
|
|
)
|
|
|
|
expected = repeated_context_mask(*inputs)
|
|
actual = repeated_context_mask(*(value.cuda() for value in inputs)).cpu()
|
|
|
|
assert torch.equal(expected, torch.tensor([True, False]))
|
|
assert torch.equal(actual, expected)
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator"
|
|
)
|
|
def test_repeated_context_mask_last_request_at_capacity():
|
|
max_model_len = 1024
|
|
all_token_ids = torch.arange(
|
|
2 * max_model_len, dtype=torch.int32, device="cuda"
|
|
).reshape(2, max_model_len)
|
|
inputs = (
|
|
all_token_ids,
|
|
torch.tensor([1], dtype=torch.int32, device="cuda"),
|
|
torch.tensor([0, 0], dtype=torch.int32, device="cuda"),
|
|
torch.tensor([max_model_len, max_model_len], dtype=torch.int32, device="cuda"),
|
|
all_token_ids[1, -4:].to(torch.int64).unsqueeze(0),
|
|
)
|
|
|
|
repeated = repeated_context_mask(*inputs)
|
|
|
|
assert not repeated.item()
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator"
|
|
)
|
|
def test_gpu_sampler_uses_fused_gumbel_for_repeated_contexts():
|
|
device = torch.device("cuda")
|
|
logits = torch.tensor([[0.0, 1.0, 2.0, 3.0], [3.0, 2.0, 1.0, 0.0]], device=device)
|
|
all_token_ids = torch.tensor(
|
|
[[101, 102, 1, 2, 1, 2], [101, 102, 3, 4, 5, 6]],
|
|
dtype=torch.int32,
|
|
device=device,
|
|
)
|
|
request_indices = torch.tensor([0, 1], dtype=torch.int64, device=device)
|
|
request_indices_np = np.array([0, 1])
|
|
temperatures = torch.ones(2, device=device)
|
|
seeds = torch.tensor([11, 22], dtype=torch.int64, device=device)
|
|
positions = torch.tensor([6, 6], dtype=torch.int64, device=device)
|
|
|
|
sampler = object.__new__(GPUWatermarkSampler)
|
|
sampler.watermarker = GumbelWatermarker(key=42, context_width=1)
|
|
sampler.deduplicate_contexts = "single_turn"
|
|
sampler.deduplicate_contexts_max_history = None
|
|
sampler.watermarking = SimpleNamespace(
|
|
np=np.ones(2, dtype=bool), gpu=torch.ones(2, dtype=torch.bool, device=device)
|
|
)
|
|
sampler.sampling_states = SimpleNamespace(
|
|
temperature=SimpleNamespace(np=np.ones(2), gpu=temperatures),
|
|
seeds=SimpleNamespace(gpu=seeds),
|
|
)
|
|
sampler.req_states = SimpleNamespace(
|
|
all_token_ids=SimpleNamespace(gpu=all_token_ids),
|
|
prompt_len=SimpleNamespace(
|
|
gpu=torch.tensor([2, 2], dtype=torch.int32, device=device)
|
|
),
|
|
total_len=SimpleNamespace(
|
|
gpu=torch.tensor([6, 6], dtype=torch.int32, device=device)
|
|
),
|
|
)
|
|
sampler.use_fp64_gumbel = False
|
|
|
|
contexts = sampler._get_contexts(request_indices)
|
|
repeated = sampler._get_repeated_contexts(request_indices, contexts)
|
|
expected = philox_gumbel_sample(
|
|
logits,
|
|
contexts,
|
|
42,
|
|
skip_mask=repeated,
|
|
expanded_idx_mapping=request_indices,
|
|
temperatures=temperatures,
|
|
seeds=seeds,
|
|
positions=positions,
|
|
)
|
|
|
|
actual, output_logits = sampler._sample_random(
|
|
logits,
|
|
request_indices,
|
|
request_indices_np,
|
|
positions,
|
|
None,
|
|
None,
|
|
False,
|
|
)
|
|
|
|
assert torch.equal(repeated, torch.tensor([True, False], device=device))
|
|
assert torch.equal(actual, expected)
|
|
assert output_logits is logits
|
|
|
|
|
|
def test_gpu_sampler_skips_watermarking_for_greedy_batch(monkeypatch):
|
|
class StubWatermarker:
|
|
context_width = 1
|
|
|
|
def sample(self, logits, contexts, random_sampler=None, skip_mask=None):
|
|
raise AssertionError("watermarker should not run for greedy requests")
|
|
|
|
sampler = object.__new__(GPUWatermarkSampler)
|
|
sampler.watermarker = StubWatermarker()
|
|
sampler.watermarking = SimpleNamespace(
|
|
np=np.array([True, True]), gpu=torch.tensor([True, True])
|
|
)
|
|
sampler.sampling_states = SimpleNamespace(
|
|
temperature=SimpleNamespace(np=np.zeros(2), gpu=torch.zeros(2)),
|
|
)
|
|
expected = (torch.tensor([3, 4]), torch.zeros(2, 8))
|
|
monkeypatch.setattr(
|
|
"vllm.v1.watermarking.gpu_sampler.Sampler._sample_random",
|
|
lambda *args, **kwargs: expected,
|
|
)
|
|
|
|
actual = sampler._sample_random(
|
|
torch.zeros(2, 8),
|
|
torch.tensor([0, 1]),
|
|
np.array([0, 1]),
|
|
torch.zeros(2, dtype=torch.int64),
|
|
None,
|
|
None,
|
|
False,
|
|
)
|
|
|
|
assert actual is expected
|
|
|
|
|
|
def test_gpu_sampler_builds_speculative_contexts_from_drafts():
|
|
sampler = object.__new__(GPUWatermarkSampler)
|
|
sampler.watermarker = SimpleNamespace(context_width=2)
|
|
sampler.req_states = SimpleNamespace(
|
|
all_token_ids=SimpleNamespace(gpu=torch.tensor([[100, 101, 10, 11, 0, 0]])),
|
|
prompt_len=SimpleNamespace(gpu=torch.tensor([2])),
|
|
total_len=SimpleNamespace(gpu=torch.tensor([4])),
|
|
)
|
|
|
|
contexts = sampler._get_contexts(
|
|
torch.tensor([0, 0, 0]),
|
|
torch.tensor([0, 1, 2]),
|
|
torch.tensor([11, 20, 21]),
|
|
)
|
|
|
|
assert torch.equal(contexts, torch.tensor([[10, 11], [11, 20], [20, 21]]))
|
|
|
|
|
|
def test_gpu_sampler_builds_chunked_multi_request_speculative_contexts():
|
|
sampler = object.__new__(GPUWatermarkSampler)
|
|
sampler.watermarker = SimpleNamespace(context_width=2)
|
|
sampler.req_states = SimpleNamespace(
|
|
all_token_ids=SimpleNamespace(
|
|
gpu=torch.tensor(
|
|
[
|
|
[100, 101, 10, 11, 0, 0],
|
|
[200, 201, 30, 31, 0, 0],
|
|
[300, 301, 50, 51, 0, 0],
|
|
]
|
|
)
|
|
),
|
|
prompt_len=SimpleNamespace(gpu=torch.tensor([2, 2, 2])),
|
|
total_len=SimpleNamespace(gpu=torch.tensor([4, 4, 4])),
|
|
)
|
|
|
|
contexts = sampler._get_contexts(
|
|
torch.tensor([2, 2, 1, 1]),
|
|
torch.tensor([0, 1, 0, 1]),
|
|
torch.tensor([51, 60, 31, 40]),
|
|
)
|
|
|
|
assert torch.equal(
|
|
contexts,
|
|
torch.tensor([[50, 51], [51, 60], [30, 31], [31, 40]]),
|
|
)
|
|
|
|
|
|
def test_draft_sampler_uses_draft_key_and_advances_context(monkeypatch):
|
|
class StubSpeculator(DraftModelSpeculator):
|
|
def capture(self): ...
|
|
|
|
def init_cudagraph_manager(self, cudagraph_mode): ...
|
|
|
|
def load_draft_model(self, target_model, target_attn_layer_names): ...
|
|
|
|
def propose(self, *args, **kwargs): ...
|
|
|
|
class StubModel:
|
|
@staticmethod
|
|
def compute_logits(hidden_states):
|
|
return torch.zeros(hidden_states.shape[0], 8)
|
|
|
|
class StubWatermarker:
|
|
@staticmethod
|
|
def sample(logits, contexts, random_sampler=None):
|
|
assert random_sampler is None
|
|
return WatermarkSample(torch.tensor([7, 7]), logits)
|
|
|
|
speculator = object.__new__(StubSpeculator)
|
|
speculator.model = StubModel()
|
|
speculator.use_fp64_gumbel = False
|
|
draft_watermarker = object.__new__(DraftWatermarker)
|
|
draft_watermarker.watermarker = StubWatermarker()
|
|
draft_watermarker.contexts = torch.tensor([[1, 2], [3, 4]])
|
|
draft_watermarker.enabled = torch.tensor([True, False])
|
|
speculator.draft_watermarker = draft_watermarker
|
|
monkeypatch.setattr(
|
|
"vllm.v1.worker.gpu.spec_decode.speculator.gumbel_sample",
|
|
lambda *args, **kwargs: torch.tensor([3, 4]),
|
|
)
|
|
|
|
sampled = speculator.sample_draft(
|
|
hidden_states=torch.zeros(2, 4),
|
|
sample_src_positions=torch.zeros(2, dtype=torch.int64),
|
|
idx_mapping=torch.tensor([0, 1]),
|
|
temperature=torch.ones(2),
|
|
seeds=torch.zeros(2, dtype=torch.int64),
|
|
draft_step=torch.tensor(0),
|
|
draft_logits=torch.zeros(2, 1, 8),
|
|
)
|
|
|
|
assert torch.equal(sampled, torch.tensor([7, 4]))
|
|
assert torch.equal(draft_watermarker.contexts, torch.tensor([[2, 7], [4, 4]]))
|
|
|
|
|
|
def test_dspark_reduced_vocab_draft_sampler_applies_watermarking(monkeypatch):
|
|
speculator = object.__new__(DSparkSpeculator)
|
|
speculator.draft_logits = torch.zeros(2, 1, 8)
|
|
speculator._d2t_scatter_index = torch.tensor([1, 5])
|
|
speculator._draft_scatter_buf = torch.full((2, 8), float("-inf"))
|
|
speculator.temperature = torch.ones(2)
|
|
speculator.seeds = torch.zeros(2, dtype=torch.int64)
|
|
speculator._step_cols = torch.tensor([0])
|
|
speculator.use_fp64_gumbel = False
|
|
watermark_logits: list[torch.Tensor] = []
|
|
|
|
def sample(logits, sampled, idx_map, temperature):
|
|
watermark_logits.append(logits.clone())
|
|
return sampled + 1
|
|
|
|
speculator.draft_watermarker = SimpleNamespace(sample=sample)
|
|
monkeypatch.setattr(
|
|
"vllm.v1.worker.gpu.spec_decode.dspark.speculator.gumbel_sample",
|
|
lambda *args, **kwargs: torch.tensor([3, 4]),
|
|
)
|
|
|
|
sampled = speculator._sample_logits(
|
|
torch.tensor([[10.0, 20.0], [30.0, 40.0]]),
|
|
torch.tensor([0, 1]),
|
|
torch.tensor([1, 1]),
|
|
0,
|
|
)
|
|
|
|
assert torch.equal(sampled, torch.tensor([4, 5]))
|
|
assert torch.equal(
|
|
watermark_logits[0][:, [1, 5]], torch.tensor([[10, 20], [30, 40]])
|
|
)
|
|
assert torch.isneginf(watermark_logits[0][:, [0, 2, 3, 4, 6, 7]]).all()
|
|
|
|
|
|
def test_dspark_target_only_watermarking_leaves_drafts_unwatermarked(monkeypatch):
|
|
speculator = object.__new__(DSparkSpeculator)
|
|
speculator.draft_logits = torch.zeros(2, 1, 8)
|
|
speculator._d2t_scatter_index = None
|
|
speculator.temperature = torch.ones(2)
|
|
speculator.seeds = torch.zeros(2, dtype=torch.int64)
|
|
speculator._step_cols = torch.tensor([0])
|
|
speculator.use_fp64_gumbel = False
|
|
speculator.draft_watermarker = None
|
|
monkeypatch.setattr(
|
|
"vllm.v1.worker.gpu.spec_decode.dspark.speculator.gumbel_sample",
|
|
lambda *args, **kwargs: torch.tensor([3, 4]),
|
|
)
|
|
|
|
sampled = speculator._sample_logits(
|
|
torch.zeros(2, 8),
|
|
torch.tensor([0, 1]),
|
|
torch.tensor([1, 1]),
|
|
0,
|
|
)
|
|
|
|
assert torch.equal(sampled, torch.tensor([3, 4]))
|