@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 0f8b398158

[watermarking] add context deduplication support to speculative decoding (#56807)

Signed-off-by: Raphael Rialland <raphael.rialland@mistral.ai> Signed-off-by: Simon Veitner <sveitner@redhat.com> Co-authored-by: Simon Veitner <sveitner@redhat.com> Co-authored-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com>

Raphaël Rialland committed Sep 30, 2026 at 18:22 UTC 0f8b398158f01fb3ea0e81b5adc93338fdf3db9c

17 files changed +1379 -88

docs/features/watermarking.md

+14 -15

@@ -77,30 +77,17 @@ watermarked. The watermark signal is diluted in proportion to the share of

77 output tokens supplied by accepted drafts; rejected drafts do not dilute it

78 because their recovery tokens are watermarked.

80 -Speculative-decoding token paths do not currently support generation-side

81 -context deduplication. The configured `deduplicate_contexts` policy is not

82 -applied to accepted drafts, rejection-recovery tokens, or bonus tokens.

83 -

80 For `dual_key_gumbel`, `alpha` has no effect under speculative decoding. The

81 speculative protocol selects the key for each token instead.

87 -## Algorithms

88 -

89 -### Gumbel-max

90 -

91 -Gumbel-max derives a deterministic pseudorandom value from the key, prior token

92 -context, and every candidate token, then uses the resulting Gumbel noise for

93 -categorical sampling. See

94 -[Aaronson's original presentation](https://simons.berkeley.edu/sites/default/files/2024-10/LLM24-2%20Slides%20-%20Scott%20Aaronson.pdf).

95 -

96 -Gumbel-max requires stochastic sampling. Greedy requests (`temperature=0`)

97 -bypass watermarking and emit a warning once per worker.

83 +## Context deduplication

85 When a token context is repeated, generation can use ordinary sampling for that

86 occurrence. Reusing the repeated context leads to a bias over the sequence, as

87 certain token choices would be correlated. Ordinary sampling at these positions

88 allows for single-sequence non-distortion (see section G.3 of the

89 [SynthID-Text supplementary materials](https://media.springernature.com/original/springer-static/esm/art%3A10.1038%2Fs41586-024-08025-4/MediaObjects/41586_2024_8025_MOESM1_ESM.pdf)).

90 +This also applies to drafter watermarking during speculative decoding.

91 The detector independently deduplicates contexts so repeated keyed random

92 vectors are not treated as independent evidence, meaning that context

93 deduplication at generation time does not reduce the watermarking signal, unless

@@ -143,6 +130,18 @@ vllm serve MODEL \

130 '{"algorithm":"gumbel","key":42,"deduplicate_contexts":"all"}'

131 ```

132

133 +## Algorithms

134 +

135 +### Gumbel-max

136 +

137 +Gumbel-max derives a deterministic pseudorandom value from the key, prior token

138 +context, and every candidate token, then uses the resulting Gumbel noise for

139 +categorical sampling. See

140 +[Aaronson's original presentation](https://simons.berkeley.edu/sites/default/files/2024-10/LLM24-2%20Slides%20-%20Scott%20Aaronson.pdf).

141 +

142 +Gumbel-max requires stochastic sampling. Greedy requests (`temperature=0`)

143 +bypass watermarking and emit a warning once per worker.

144 +

145 ### Dual-key Gumbel-max

146

147 Dual-key Gumbel-max derives independent keys A and B from one configured master

tests/test_config.py

+23 -2

@@ -3526,7 +3526,6 @@ def test_target_only_gumbel_allows_speculative_decoding(caplog_vllm, disable_log

3526 config._check_watermarking_unsupported()

3527

3528 assert "Target-only watermarking leaves accepted draft tokens" in caplog_vllm.text

3529 - assert "Context deduplication is not supported" in caplog_vllm.text

3529

3530

3531 def test_speculative_watermarking_without_context_dedup_does_not_warn(

@@ -3543,10 +3542,32 @@ def test_speculative_watermarking_without_context_dedup_does_not_warn(

3542 parallel_drafting=False,

3543 )

3544

3545 + caplog_vllm.clear()

3546 with caplog_vllm.at_level(logging.WARNING):

3547 config._check_watermarking_unsupported()

3548

3549 - assert "Context deduplication is not supported" not in caplog_vllm.text

3549 + assert "dedup" not in caplog_vllm.text.lower()

3550 +

3551 +

3552 +def test_speculative_context_dedup_logs_no_unsupported_warning(

3553 + caplog_vllm, disable_log_dedup

3554 +):

3555 + config = _watermarked_vllm_config()

3556 + config.watermark_config = WatermarkConfig(

3557 + algorithm="dual_key_gumbel", key=42, deduplicate_contexts="single_turn"

3558 + )

3559 + config.speculative_config = SimpleNamespace(

3560 + method="mtp",

3561 + draft_sample_method="probabilistic",

3562 + rejection_sample_method="standard",

3563 + parallel_drafting=False,

3564 + )

3565 +

3566 + caplog_vllm.clear()

3567 + with caplog_vllm.at_level(logging.WARNING):

3568 + config._check_watermarking_unsupported()

3569 +

3570 + assert "dedup" not in caplog_vllm.text.lower()

3571

3572

3573 def test_gumbel_rejects_speculative_decoding_without_target_only():

tests/v1/spec_decode/test_rejection_sampler_utils.py

+176

@@ -203,6 +203,182 @@ def test_watermarked_recovery_uses_target_key():

203 assert sampled[0, 0] == expected.token_ids[0]

204

205

206 +def test_repeated_draft_context_does_not_change_acceptance():

207 + torch.manual_seed(0)

208 + num_trials = 8

209 + # A draft distribution of its own, so the drafts are not accepted outright.

210 + target_logits = torch.randn(64, device="cuda")

211 + draft_logits = torch.randn(64, device="cuda")

212 + inputs = _build_rejection_sample_inputs(

213 + target_logits,

214 + draft_logits,

215 + num_speculative_steps=1,

216 + temperature=1.0,

217 + num_trials=num_trials,

218 + )

219 + expected, expected_num_sampled = rejection_sample(**inputs, num_speculative_steps=1)

220 +

221 + sampled, num_sampled = watermarked_rejection_sample(

222 + **inputs,

223 + num_speculative_steps=1,

224 + contexts=torch.tensor([[10, 11], [11, 12]], device="cuda").repeat(

225 + num_trials, 1

226 + ),

227 + watermarking=torch.ones(num_trials, dtype=torch.bool, device="cuda"),

228 + watermarking_skip_mask=torch.tensor([True, False], device="cuda").repeat(

229 + num_trials

230 + ),

231 + watermarker=GumbelWatermarker(key=42, context_width=2),

232 + )

233 +

234 + # Some drafts must actually be rejected, or the ordinary-residual branch

235 + # the skip mask selects is never reached.

236 + assert (expected_num_sampled == 1).any()

237 + torch.testing.assert_close(num_sampled, expected_num_sampled, rtol=0, atol=0)

238 + torch.testing.assert_close(sampled[:, :1], expected[:, :1], rtol=0, atol=0)

239 +

240 +

241 +def test_repeated_context_skip_mask_is_read_per_row_in_a_block():

242 + """A block of three drafts, accepted at row 0 and rejected at row 1.

243 +

244 + The skip mask selects the sampler per row, so the row the target recovers

245 + at decides which key draws the recovery token.

246 + """

247 + num_speculative_steps = 3

248 + target_logits = torch.tensor([float("-inf"), 0.0, 1.0, 2.0], device="cuda")

249 + draft_logits = torch.tensor([30.0, -30.0, -30.0, -30.0], device="cuda")

250 + inputs = _build_rejection_sample_inputs(

251 + target_logits,

252 + draft_logits,

253 + num_speculative_steps=num_speculative_steps,

254 + temperature=1.0,

255 + num_trials=1,

256 + )

257 + # Row 0 drafts a token the target gives 0.67 and the draft e**-60, so it is

258 + # accepted; row 1 drafts a token the target gives no mass, so it is not.

259 + inputs["draft_sampled"][1] = 3

260 + inputs["draft_sampled"][2] = 0

261 + contexts = torch.tensor(

262 + [[10, 11], [11, 3], [3, 0], [0, 1]], dtype=torch.int64, device="cuda"

263 + )

264 + expected, expected_num_sampled = rejection_sample(

265 + **inputs, num_speculative_steps=num_speculative_steps

266 + )

267 + assert expected_num_sampled.item() == 2

268 + assert expected[0, 0].item() == 3

269 +

270 + watermarker = GumbelWatermarker(key=42, context_width=2)

271 + skipped, skipped_num_sampled = watermarked_rejection_sample(

272 + **inputs,

273 + num_speculative_steps=num_speculative_steps,

274 + contexts=contexts,

275 + watermarking=torch.tensor([True], device="cuda"),

276 + watermarking_skip_mask=torch.tensor([False, True, False, False], device="cuda"),

277 + watermarker=watermarker,

278 + )

279 +

280 + assert skipped_num_sampled.item() == 2

281 + torch.testing.assert_close(skipped[:, :2], expected[:, :2], rtol=0, atol=0)

282 +

283 + watermarked, watermarked_num_sampled = watermarked_rejection_sample(

284 + **inputs,

285 + num_speculative_steps=num_speculative_steps,

286 + contexts=contexts,

287 + watermarking=torch.tensor([True], device="cuda"),

288 + watermarking_skip_mask=torch.zeros(

289 + num_speculative_steps + 1, dtype=torch.bool, device="cuda"

290 + ),

291 + watermarker=watermarker,

292 + )

293 + # The residual equals the target to within e**-60, so the recovery token is

294 + # the target key's draw on the rejected row's context.

295 + recovery = GumbelWatermarker(key=42, context_width=2).sample(

296 + target_logits.unsqueeze(0), contexts[1:2], lambda _: None

297 + )

298 +

299 + assert watermarked_num_sampled.item() == 2

300 + assert watermarked[0, 0].item() == 3

301 + assert watermarked[0, 1] == recovery.token_ids[0]

302 +

303 +

304 +def test_repeated_recovery_context_uses_ordinary_residual_sample():

305 + target_logits = torch.tensor([float("-inf"), 0.0, 1.0, 2.0], device="cuda")

306 + draft_logits = torch.tensor([100.0, -100.0, -100.0, -100.0], device="cuda")

307 + inputs = _build_rejection_sample_inputs(

308 + target_logits,

309 + draft_logits,

310 + num_speculative_steps=1,

311 + temperature=1.0,

312 + num_trials=1,

313 + )

314 + inputs["draft_sampled"][1] = 0

315 + expected, expected_num_sampled = rejection_sample(**inputs, num_speculative_steps=1)

316 +

317 + sampled, num_sampled = watermarked_rejection_sample(

318 + **inputs,

319 + num_speculative_steps=1,

320 + contexts=torch.tensor([[10, 11], [11, 0]], device="cuda"),

321 + watermarking=torch.tensor([True], device="cuda"),

322 + watermarking_skip_mask=torch.tensor([True, False], device="cuda"),

323 + watermarker=GumbelWatermarker(key=42, context_width=2),

324 + )

325 +

326 + torch.testing.assert_close(num_sampled, expected_num_sampled, rtol=0, atol=0)

327 + torch.testing.assert_close(sampled[:, :1], expected[:, :1], rtol=0, atol=0)

328 +

329 +

330 +def test_repeated_bonus_context_uses_ordinary_target_sample():

331 + target_logits = torch.randn(64, device="cuda")

332 + inputs = _build_rejection_sample_inputs(

333 + target_logits,

334 + target_logits,

335 + num_speculative_steps=1,

336 + temperature=1.0,

337 + num_trials=1,

338 + )

339 + expected, expected_num_sampled = rejection_sample(**inputs, num_speculative_steps=1)

340 +

341 + sampled, num_sampled = watermarked_rejection_sample(

342 + **inputs,

343 + num_speculative_steps=1,

344 + contexts=torch.tensor([[10, 11], [11, 12]], device="cuda"),

345 + watermarking=torch.tensor([True], device="cuda"),

346 + watermarking_skip_mask=torch.tensor([False, True], device="cuda"),

347 + watermarker=GumbelWatermarker(key=42, context_width=2),

348 + )

349 +

350 + torch.testing.assert_close(num_sampled, expected_num_sampled, rtol=0, atol=0)

351 + torch.testing.assert_close(sampled[:, :2], expected[:, :2], rtol=0, atol=0)

352 +

353 +

354 +def test_repeated_context_does_not_change_unwatermarked_or_greedy_rows():

355 + target_logits = torch.randn(64, device="cuda")

356 + inputs = _build_rejection_sample_inputs(

357 + target_logits,

358 + target_logits,

359 + num_speculative_steps=1,

360 + temperature=1.0,

361 + num_trials=2,

362 + )

363 + inputs["temperature"][1] = 0

364 + expected, expected_num_sampled = rejection_sample(**inputs, num_speculative_steps=1)

365 +

366 + sampled, num_sampled = watermarked_rejection_sample(

367 + **inputs,

368 + num_speculative_steps=1,

369 + contexts=torch.tensor([[10, 11], [11, 12], [20, 21], [21, 22]], device="cuda"),

370 + watermarking=torch.tensor([False, True], device="cuda"),

371 + watermarking_skip_mask=torch.ones(4, dtype=torch.bool, device="cuda"),

372 + watermarker=GumbelWatermarker(key=42, context_width=2),

373 + )

374 +

375 + torch.testing.assert_close(num_sampled, expected_num_sampled, rtol=0, atol=0)

376 + for row, length in enumerate(num_sampled):

377 + torch.testing.assert_close(

378 + sampled[row, : int(length)], expected[row, : int(length)], rtol=0, atol=0

379 + )

380 +

381 +

382 def test_watermarked_recovery_supports_smaller_draft_vocabulary():

383 target_logits = torch.tensor([float("-inf"), 0.0, 1.0, 2.0, 100.0], device="cuda")

384 draft_logits = torch.tensor([100.0, -100.0, -100.0, -100.0], device="cuda")

tests/watermarking/test_gumbel.py

+62

@@ -12,6 +12,7 @@ from vllm.v1.watermarking import (

12 derive_watermark_key,

13 )

14 from vllm.v1.watermarking.gumbel import _gamma_survival_integer_shape

15 +from vllm.v1.watermarking.watermarker import RandomSampler

16 from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample

17 from vllm.v1.worker.gpu.sample.watermark import philox_gumbel_sample

@@ -197,3 +198,64 @@ def test_philox_gumbel_sample_skip_mask_matches_separate_samplers(use_fp64: bool

198 )

199

200 torch.testing.assert_close(actual, expected, rtol=0, atol=0)

201 +

202 +

203 +@pytest.mark.skipif(

204 + not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator"

205 +)

206 +def test_philox_mixed_draft_sampling_caches_unprocessed_logits():

207 + torch.manual_seed(0)

208 + num_rows, vocab_size, num_steps = 5, 4099, 3

209 + logits = torch.randn(num_rows, vocab_size, device="cuda", dtype=torch.float16)

210 + temperatures = torch.tensor([0.0, 0.3, 0.5, 1.0, 2.0], device="cuda")

211 + processed_logits = logits / torch.where(

212 + temperatures == 0, 1, temperatures

213 + ).unsqueeze(-1)

214 + contexts = torch.randint(0, vocab_size, (num_rows, 4), device="cuda")

215 + idx_mapping = torch.arange(num_rows, dtype=torch.int32, device="cuda")

216 + seeds = torch.arange(num_rows, dtype=torch.int64, device="cuda") + 10

217 + positions = torch.arange(num_rows, dtype=torch.int64, device="cuda") + 20

218 + # The skipped rows must divide by a temperature other than 1, or the

219 + # unprocessed logits the cache stores are the logits the sampler sees.

220 + # Row 0 keeps a greedy row in the skipped set.

221 + skip_mask = torch.tensor([True, True, False, False, True], device="cuda")

222 + cols = torch.arange(num_rows, dtype=torch.int32, device="cuda") % num_steps

223 + cache = torch.zeros(

224 + num_rows, num_steps, vocab_size + 1, dtype=logits.dtype, device="cuda"

225 + )

226 +

227 + ordinary = gumbel_sample(

228 + logits,

229 + idx_mapping,

230 + temperatures,

231 + seeds,

232 + positions,

233 + apply_temperature=True,

234 + is_drafting=True,

235 + )

236 + watermarked = philox_gumbel_sample(processed_logits, contexts, 42)

237 + expected = torch.where(skip_mask, ordinary, watermarked)

238 + actual = (

239 + GumbelWatermarker(key=42, context_width=4)

240 + .sample(

241 + processed_logits,

242 + contexts,

243 + random_sampler=RandomSampler(

244 + expanded_idx_mapping=idx_mapping,

245 + temperatures=temperatures,

246 + seeds=seeds,

247 + positions=positions,

248 + is_drafting=True,

249 + logits_cache=cache,

250 + logits_cache_col=cols,

251 + logits_cache_source=logits,

252 + ),

253 + skip_mask=skip_mask,

254 + )

255 + .token_ids

256 + )

257 +

258 + torch.testing.assert_close(actual, expected, rtol=0, atol=0)

259 + stored = cache[torch.arange(num_rows, device="cuda"), cols.long(), :vocab_size]

260 + assert torch.equal(stored.view(torch.int16), logits.view(torch.int16))

261 + assert not cache[:, :, -1].any()

tests/watermarking/test_watermarking.py

+550 -13

@@ -28,6 +28,7 @@ from vllm.v1.watermarking.spec_decode import (

28 from vllm.v1.watermarking.watermarker import Watermarker, WatermarkSample

29 from vllm.v1.worker.gpu.sample.sampler import Sampler

30 from vllm.v1.worker.gpu.sample.watermark import (

31 + draft_watermarking_mask,

32 philox_gumbel_sample,

33 repeated_context_mask,

34 )

@@ -202,11 +203,17 @@ def test_speculative_decoding_uses_fixed_dual_key_roles():

203 max_num_reqs=1,

204 device=torch.device("cpu"),

205 allow_target_only=False,

206 + num_speculative_steps=3,

207 + deduplicate_contexts="all",

208 + deduplicate_contexts_max_history=None,

209 )

210

211 assert target.prf.key == derive_watermark_key(42, b"key_b")

212 assert draft is not None

213 assert draft.watermarker.prf.key == derive_watermark_key(42, b"key_a")

214 + assert draft.num_speculative_steps == 3

215 + assert draft.deduplicate_contexts == "all"

216 + assert draft.deduplicate_contexts_max_history is None

217

218

219 def test_recovery_key_rejects_the_unsplit_dual_key_watermarker():

@@ -892,6 +899,187 @@ def test_gpu_sampler_builds_speculative_contexts_from_drafts():

899 assert torch.equal(contexts, torch.tensor([[10, 11], [11, 20], [20, 21]]))

900

901

902 +def test_gpu_sampler_deduplicates_contexts_within_speculative_block():

903 + sampler = object.__new__(GPUWatermarkSampler)

904 + sampler.watermarker = SimpleNamespace(context_width=2)

905 + sampler.deduplicate_contexts = "single_turn"

906 + sampler.deduplicate_contexts_max_history = None

907 + sampler.num_speculative_tokens = 4

908 + sampler.req_states = SimpleNamespace(

909 + all_token_ids=SimpleNamespace(gpu=torch.tensor([[100, 1, 2, 0, 0, 0]])),

910 + prompt_len=SimpleNamespace(gpu=torch.tensor([1])),

911 + total_len=SimpleNamespace(gpu=torch.tensor([3])),

912 + )

913 + request_indices = torch.tensor([0, 0, 0, 0, 0])

914 + local_positions = torch.tensor([0, 1, 2, 3, 4])

915 + contexts = sampler._get_contexts(

916 + request_indices,

917 + local_positions,

918 + torch.tensor([0, 3, 4, 1, 2]),

919 + )

920 +

921 + repeated = sampler._get_repeated_contexts(

922 + request_indices, contexts, local_positions

923 + )

924 +

925 + assert torch.equal(

926 + contexts,

927 + torch.tensor([[1, 2], [2, 3], [3, 4], [4, 1], [1, 2]]),

928 + )

929 + assert torch.equal(repeated, torch.tensor([False, False, False, False, True]))

930 +

931 + sampler.deduplicate_contexts_max_history = 3

932 + repeated = sampler._get_repeated_contexts(

933 + request_indices, contexts, local_positions

934 + )

935 + assert not repeated.any()

936 +

937 +

938 +def test_gpu_sampler_all_scope_skips_only_partial_speculative_contexts():

939 + sampler = object.__new__(GPUWatermarkSampler)

940 + sampler.watermarker = SimpleNamespace(context_width=2)

941 + sampler.deduplicate_contexts = "all"

942 + sampler.deduplicate_contexts_max_history = None

943 + sampler.num_speculative_tokens = 3

944 + sampler.req_states = SimpleNamespace(

945 + all_token_ids=SimpleNamespace(gpu=torch.tensor([[100, 101, 10, 0, 0]])),

946 + prompt_len=SimpleNamespace(gpu=torch.tensor([2])),

947 + total_len=SimpleNamespace(gpu=torch.tensor([3])),

948 + )

949 + request_indices = torch.tensor([0, 0, 0])

950 + local_positions = torch.tensor([0, 1, 2])

951 + contexts = sampler._get_contexts(

952 + request_indices,

953 + local_positions,

954 + torch.tensor([10, 20, 30]),

955 + )

956 +

957 + repeated = sampler._get_repeated_contexts(

958 + request_indices, contexts, local_positions

959 + )

960 +

961 + assert torch.equal(contexts, torch.tensor([[-1, 10], [10, 20], [20, 30]]))

962 + assert torch.equal(repeated, torch.tensor([True, False, False]))

963 +

964 +

965 +@pytest.mark.skipif(

966 + not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator"

967 +)

968 +@pytest.mark.parametrize("max_history,expected_last", [(None, True), (3, False)])

969 +def test_speculative_context_repetition_accelerator_parity(

970 + max_history: int | None, expected_last: bool

971 +):

972 + all_token_ids = torch.tensor([[10, 11]])

973 + req_indices = torch.tensor([0, 0, 0, 0, 0])

974 + prompt_lens = torch.tensor([0])

975 + total_lens = torch.tensor([2])

976 + local_positions = torch.tensor([0, 1, 2, 3, 4])

977 + contexts = torch.tensor([[1, 2], [2, 3], [3, 4], [4, 1], [1, 2]])

978 +

979 + expected = repeated_context_mask(

980 + all_token_ids,

981 + req_indices,

982 + prompt_lens,

983 + total_lens,

984 + contexts,

985 + history_offsets=local_positions,

986 + local_positions=local_positions,

987 + num_speculative_steps=4,

988 + max_history=max_history,

989 + )

990 + actual = repeated_context_mask(

991 + all_token_ids.cuda(),

992 + req_indices.cuda(),

993 + prompt_lens.cuda(),

994 + total_lens.cuda(),

995 + contexts.cuda(),

996 + history_offsets=local_positions.cuda(),

997 + local_positions=local_positions.cuda(),

998 + num_speculative_steps=4,

999 + max_history=max_history,

1000 + ).cpu()

1001 +

1002 + assert torch.equal(

1003 + expected, torch.tensor([False, False, False, False, expected_last])

1004 + )

1005 + assert torch.equal(actual, expected)

1006 +

1007 +

1008 +@pytest.mark.skipif(

1009 + not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator"

1010 +)

1011 +@pytest.mark.parametrize("steps_layout", ["per_row", "broadcast"])

1012 +@pytest.mark.parametrize("req_indices_layout", ["contiguous", "column_view"])

1013 +def test_draft_context_deduplication_accelerator_parity(

1014 + steps_layout, req_indices_layout

1015 +):

1016 + """Parity with the CPU reference for the layouts the speculators pass.

1017 +

1018 + A draft step is one 0-d tensor broadcast over the batch, and DSpark samples

1019 + a step at a time, so its request indices are an `idx_map[:, i]` column. Both

1020 + are read per row, so a row that ignores the layout consults another row's

1021 + step or another request's history.

1022 + """

1023 + num_rows, num_steps, step = 4, 3, 1

1024 + all_token_ids = torch.tensor([[10, 11, 20, 11], [30, 31, 11, 20]])

1025 + prompt_lens = torch.tensor([0, 0])

1026 + total_lens = torch.tensor([4, 4])

1027 + contexts = torch.tensor([[40, 41], [31, 11], [10, 11], [50, 51]])

1028 + prior_contexts = torch.zeros(num_rows, num_steps, 2, dtype=torch.int64)

1029 + prior_contexts[0, 0] = torch.tensor([40, 41])

1030 + enabled = torch.ones(num_rows, dtype=torch.bool)

1031 + req_indices = torch.tensor([0, 1, 0, 1])

1032 +

1033 + def run(all_token_ids, req_indices, steps, prior):

1034 + return draft_watermarking_mask(

1035 + all_token_ids,

1036 + req_indices,

1037 + prompt_lens.to(all_token_ids.device),

1038 + total_lens.to(all_token_ids.device),

1039 + prior,

1040 + contexts.to(all_token_ids.device),

1041 + steps,

1042 + enabled.to(all_token_ids.device),

1043 + max_history=3,

1044 + include_prompt=False,

1045 + )

1046 +

1047 + reference_prior = prior_contexts.clone()

1048 + reference = run(

1049 + all_token_ids,

1050 + req_indices,

1051 + torch.full((num_rows,), step),

1052 + reference_prior,

1053 + )

1054 +

1055 + # Row 0 repeats its own step-0 context, rows 1 and 2 repeat a committed

1056 + # context of their own request, row 3 repeats nothing.

1057 + assert torch.equal(reference, torch.tensor([False, False, False, True]))

1058 + expected_prior = prior_contexts.clone()

1059 + expected_prior[:, step] = contexts

1060 + assert torch.equal(reference_prior, expected_prior)

1061 +

1062 + if steps_layout == "per_row":

1063 + steps = torch.full((num_rows,), step).cuda()

1064 + else:

1065 + steps = torch.tensor(step).cuda().expand(num_rows)

1066 + assert steps.stride(0) == 0

1067 +

1068 + if req_indices_layout == "contiguous":

1069 + gpu_req_indices = req_indices.cuda()

1070 + else:

1071 + gpu_req_indices = (

1072 + req_indices.repeat_interleave(num_steps).view(num_rows, num_steps).cuda()

1073 + )[:, 1]

1074 + assert gpu_req_indices.stride(0) == num_steps

1075 +

1076 + actual_prior = prior_contexts.cuda()

1077 + actual = run(all_token_ids.cuda(), gpu_req_indices, steps, actual_prior).cpu()

1078 +

1079 + assert torch.equal(actual, reference)

1080 + assert torch.equal(actual_prior.cpu(), expected_prior)

1081 +

1082 +

1083 def test_gpu_sampler_builds_chunked_multi_request_speculative_contexts():

1084 sampler = object.__new__(GPUWatermarkSampler)

1085 sampler.watermarker = SimpleNamespace(context_width=2)

@@ -921,6 +1109,69 @@ def test_gpu_sampler_builds_chunked_multi_request_speculative_contexts():

1109 )

1110

1111

1112 +@pytest.mark.skipif(

1113 + not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator"

1114 +)

1115 +def test_draft_context_deduplication_broadcasts_the_draft_step():

1116 + """The draft step is one 0-d tensor broadcast over the batch.

1117 +

1118 + The context occurs four positions back, which `max_history - step` puts

1119 + just outside the scanned window, so a row that mistakes another row's step

1120 + for its own reports the opposite verdict. The neighbours leg is the

1121 + deterministic control: a row that reads a step of 0 instead of 2 widens its

1122 + window by two positions and finds the context. The production leg asserts

1123 + the layout the speculators pass, where the stride-less read lands outside

1124 + the allocation and can return the right value by chance.

1125 + """

1126 + num_rows = 4

1127 + all_token_ids = torch.tensor([[3, 4, 9, 9, 9, 9], [3, 4, 9, 9, 9, 9]])

1128 + req_indices = torch.tensor([0, 0, 1, 1])

1129 + prompt_lens = torch.tensor([0, 0])

1130 + total_lens = torch.tensor([6, 6])

1131 + prior_contexts = torch.zeros(num_rows, 3, 2, dtype=torch.int64)

1132 + contexts = torch.tensor([[3, 4]]).expand(num_rows, 2).contiguous()

1133 + enabled = torch.ones(num_rows, dtype=torch.bool)

1134 +

1135 + def run(steps, device):

1136 + prior = prior_contexts.clone().to(device)

1137 + active = draft_watermarking_mask(

1138 + all_token_ids.to(device),

1139 + req_indices.to(device),

1140 + prompt_lens.to(device),

1141 + total_lens.to(device),

1142 + prior,

1143 + contexts.to(device),

1144 + steps,

1145 + enabled.to(device),

1146 + max_history=4,

1147 + include_prompt=False,

1148 + )

1149 + return active.cpu(), prior.cpu()

1150 +

1151 + expected, expected_prior = run(torch.full((num_rows,), 2), torch.device("cpu"))

1152 + assert torch.equal(expected, torch.ones(num_rows, dtype=torch.bool))

1153 +

1154 + materialized, materialized_prior = run(

1155 + torch.full((num_rows,), 2).cuda(), torch.device("cuda")

1156 + )

1157 + assert torch.equal(materialized, expected)

1158 + assert torch.equal(materialized_prior, expected_prior)

1159 +

1160 + # Neighbouring steps are the value a row reads when the broadcast stride is

1161 + # ignored; the production tensor reads uninitialized memory instead.

1162 + neighbours = torch.tensor([2, 0, 0, 0]).cuda()

1163 + broadcast, broadcast_prior = run(

1164 + neighbours[0].expand(num_rows), torch.device("cuda")

1165 + )

1166 + assert torch.equal(broadcast, expected)

1167 + assert torch.equal(broadcast_prior, expected_prior)

1168 +

1169 + step = torch.tensor(2).cuda()

1170 + production, production_prior = run(step.expand(num_rows), torch.device("cuda"))

1171 + assert torch.equal(production, expected)

1172 + assert torch.equal(production_prior, expected_prior)

1173 +

1174 +

1175 def test_draft_sampler_uses_draft_key_and_advances_context(monkeypatch):

1176 class StubSpeculator(DraftModelSpeculator):

1177 def capture(self): ...

@@ -937,22 +1188,32 @@ def test_draft_sampler_uses_draft_key_and_advances_context(monkeypatch):

1188 return torch.zeros(hidden_states.shape[0], 8)

1189

1190 class StubWatermarker:

1191 + context_width = 2

1192 +

1193 @staticmethod

941 - def sample(logits, contexts, random_sampler=None):

942 - assert random_sampler is None

943 - return WatermarkSample(torch.tensor([7, 7]), logits)

1194 + def sample(logits, contexts, random_sampler=None, skip_mask=None):

1195 + assert random_sampler is not None

1196 + return WatermarkSample(

1197 + torch.where(skip_mask, random_sampler(logits), 7), logits

1198 + )

1199

1200 speculator = object.__new__(StubSpeculator)

1201 speculator.model = StubModel()

1202 speculator.use_fp64_gumbel = False

1203 speculator.acceptance_estimator = None

949 - draft_watermarker = object.__new__(DraftWatermarker)

950 - draft_watermarker.watermarker = StubWatermarker()

951 - draft_watermarker.contexts = torch.tensor([[1, 2], [3, 4]])

1204 + draft_watermarker = DraftWatermarker(

1205 + StubWatermarker(),

1206 + max_num_reqs=2,

1207 + device=torch.device("cpu"),

1208 + num_speculative_steps=1,

1209 + deduplicate_contexts="none",

1210 + deduplicate_contexts_max_history=None,

1211 + )

1212 + draft_watermarker.contexts.copy_(torch.tensor([[1, 2], [3, 4]]))

1213 draft_watermarker.enabled = torch.tensor([True, False])

1214 speculator.draft_watermarker = draft_watermarker

1215 monkeypatch.setattr(

955 - "vllm.v1.worker.gpu.spec_decode.speculator.gumbel_sample",

1216 + "vllm.v1.watermarking.watermarker.gumbel_sample",

1217 lambda *args, **kwargs: torch.tensor([3, 4]),

1218 )

1219

@@ -970,6 +1231,274 @@ def test_draft_sampler_uses_draft_key_and_advances_context(monkeypatch):

1231 assert torch.equal(draft_watermarker.contexts, torch.tensor([[2, 7], [4, 4]]))

1232

1233

1234 +def test_draft_sampler_uses_ordinary_samples_for_repeated_contexts(monkeypatch):

1235 + class StubWatermarker:

1236 + context_width = 1

1237 +

1238 + def __init__(self):

1239 + self.samples = iter((2, 1, 8))

1240 +

1241 + def sample(self, logits, contexts, random_sampler=None, skip_mask=None):

1242 + watermarked = torch.tensor([next(self.samples)])

1243 + return WatermarkSample(

1244 + torch.where(skip_mask, random_sampler(logits), watermarked), logits

1245 + )

1246 +

1247 + draft_watermarker = DraftWatermarker(

1248 + StubWatermarker(),

1249 + max_num_reqs=1,

1250 + device=torch.device("cpu"),

1251 + num_speculative_steps=3,

1252 + deduplicate_contexts="single_turn",

1253 + deduplicate_contexts_max_history=None,

1254 + )

1255 + draft_watermarker.prepare(

1256 + contexts=torch.tensor([[1]]),

1257 + enabled=torch.tensor([True]),

1258 + all_token_ids=torch.tensor([[100, 1, 0, 0]]),

1259 + prompt_lens=torch.tensor([1]),

1260 + total_lens=torch.tensor([2]),

1261 + )

1262 + idx_mapping = torch.tensor([0])

1263 + temperature = torch.tensor([1.0])

1264 + ordinary_samples = iter((torch.tensor([5]), torch.tensor([6]), torch.tensor([4])))

1265 + monkeypatch.setattr(

1266 + "vllm.v1.watermarking.watermarker.gumbel_sample",

1267 + lambda *args, **kwargs: next(ordinary_samples),

1268 + )

1269 + seeds = torch.tensor([0])

1270 + draft_logits = torch.zeros(1, 3, 8)

1271 +

1272 + first = draft_watermarker.sample(

1273 + torch.zeros(1, 8),

1274 + idx_mapping,

1275 + temperature=temperature,

1276 + seed=seeds,

1277 + pos=torch.tensor([0]),

1278 + apply_temperature=True,

1279 + is_drafting=True,

1280 + logits_cache_col=torch.tensor(0),

1281 + logits_cache=draft_logits,

1282 + use_fp64=False,

1283 + )

1284 + second = draft_watermarker.sample(

1285 + torch.zeros(1, 8),

1286 + idx_mapping,

1287 + temperature=temperature,

1288 + seed=seeds,

1289 + pos=torch.tensor([1]),

1290 + apply_temperature=True,

1291 + is_drafting=True,

1292 + logits_cache_col=torch.tensor(1),

1293 + logits_cache=draft_logits,

1294 + use_fp64=False,

1295 + )

1296 + repeated = draft_watermarker.sample(

1297 + torch.zeros(1, 8),

1298 + idx_mapping,

1299 + temperature=temperature,

1300 + seed=seeds,

1301 + pos=torch.tensor([2]),

1302 + apply_temperature=True,

1303 + is_drafting=True,

1304 + logits_cache_col=torch.tensor(2),

1305 + logits_cache=draft_logits,

1306 + use_fp64=False,

1307 + )

1308 +

1309 + assert first.item() == 2

1310 + assert second.item() == 1

1311 + assert repeated.item() == 4

1312 +

1313 +

1314 +@pytest.mark.skipif(

1315 + not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator"

1316 +)

1317 +def test_draft_context_deduplication_reads_strided_request_indices():

1318 + """DSpark samples a step at a time, so req_indices is an `idx_map[:, i]` column."""

1319 + num_reqs = 4

1320 + num_speculative_steps = 3

1321 + all_token_ids = torch.tensor(

1322 + [[3, 4, 3, 4], [9, 9, 9, 9], [3, 4, 3, 4], [9, 9, 9, 9]]

1323 + ).cuda()

1324 + prompt_lens = torch.zeros(num_reqs, dtype=torch.int64).cuda()

1325 + total_lens = torch.full((num_reqs,), 4).cuda()

1326 + contexts = torch.tensor([[3, 4]]).expand(num_reqs, 2).contiguous().cuda()

1327 + steps = torch.zeros(num_reqs, dtype=torch.int64).cuda()

1328 + enabled = torch.ones(num_reqs, dtype=torch.bool).cuda()

1329 +

1330 + idx_map = (

1331 + torch.arange(num_reqs)

1332 + .repeat_interleave(num_speculative_steps)

1333 + .view(num_reqs, num_speculative_steps)

1334 + .cuda()

1335 + )

1336 + column = idx_map[:, 1]

1337 + assert column.stride(0) == num_speculative_steps

1338 +

1339 + def run(req_indices):

1340 + prior = torch.zeros(

1341 + num_reqs, num_speculative_steps, 2, dtype=torch.int64, device="cuda"

1342 + )

1343 + return draft_watermarking_mask(

1344 + all_token_ids,

1345 + req_indices,

1346 + prompt_lens,

1347 + total_lens,

1348 + prior,

1349 + contexts,

1350 + steps,

1351 + enabled,

1352 + max_history=None,

1353 + include_prompt=False,

1354 + ).cpu()

1355 +

1356 + expected = torch.tensor([False, True, False, True])

1357 + assert torch.equal(run(column.contiguous()), expected)

1358 + assert torch.equal(run(column), expected)

1359 +

1360 +

1361 +def test_draft_sampler_deduplicates_against_committed_history(monkeypatch):

1362 + class StubWatermarker:

1363 + context_width = 1

1364 +

1365 + @staticmethod

1366 + def sample(logits, contexts, random_sampler=None, skip_mask=None):

1367 + return WatermarkSample(

1368 + torch.where(skip_mask, random_sampler(logits), 7), logits

1369 + )

1370 +

1371 + draft_watermarker = DraftWatermarker(

1372 + StubWatermarker(),

1373 + max_num_reqs=1,

1374 + device=torch.device("cpu"),

1375 + num_speculative_steps=1,

1376 + deduplicate_contexts="single_turn",

1377 + deduplicate_contexts_max_history=None,

1378 + )

1379 + draft_watermarker.prepare(

1380 + contexts=torch.tensor([[1]]),

1381 + enabled=torch.tensor([True]),

1382 + all_token_ids=torch.tensor([[100, 1, 2, 1]]),

1383 + prompt_lens=torch.tensor([1]),

1384 + total_lens=torch.tensor([4]),

1385 + )

1386 +

1387 + monkeypatch.setattr(

1388 + "vllm.v1.watermarking.watermarker.gumbel_sample",

1389 + lambda *args, **kwargs: torch.tensor([4]),

1390 + )

1391 + sampled = draft_watermarker.sample(

1392 + torch.zeros(1, 8),

1393 + torch.tensor([0]),

1394 + temperature=torch.tensor([1.0]),

1395 + seed=torch.tensor([0]),

1396 + pos=torch.tensor([0]),

1397 + apply_temperature=True,

1398 + is_drafting=True,

1399 + logits_cache_col=torch.tensor(0),

1400 + logits_cache=torch.zeros(1, 1, 8),

1401 + use_fp64=False,

1402 + )

1403 +

1404 + assert sampled.item() == 4

1405 +

1406 +

1407 +@pytest.mark.skipif(

1408 + not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator"

1409 +)

1410 +@pytest.mark.parametrize("idx_mapping_layout", ["contiguous", "column_view"])

1411 +def test_draft_watermarker_deduplicates_across_a_block_on_cuda(

1412 + monkeypatch, idx_mapping_layout

1413 +):

1414 + """Drive the whole draft sampler over a block with the production shapes.

1415 +

1416 + `draft_step` is a 0-d tensor, and a speculator that samples one step at a

1417 + time hands over a column of its index map. Each step is checked against the

1418 + committed history the window still reaches and against the earlier steps of

1419 + the block.

1420 + """

1421 + num_reqs, num_steps, vocab_size = 2, 3, 16

1422 + watermarker = GumbelWatermarker(key=42, context_width=2)

1423 + draft_watermarker = DraftWatermarker(

1424 + watermarker,

1425 + max_num_reqs=num_reqs,

1426 + device=torch.device("cuda"),

1427 + num_speculative_steps=num_steps,

1428 + deduplicate_contexts="single_turn",

1429 + deduplicate_contexts_max_history=4,

1430 + )

1431 + draft_watermarker.prepare(

1432 + contexts=torch.tensor([[5, 6], [6, 3]]).cuda(),

1433 + enabled=torch.ones(num_reqs, dtype=torch.bool).cuda(),

1434 + all_token_ids=torch.tensor([[5, 6, 7, 5, 6], [8, 9, 6, 6, 3]]).cuda(),

1435 + prompt_lens=torch.zeros(num_reqs, dtype=torch.int64).cuda(),

1436 + total_lens=torch.full((num_reqs,), 5).cuda(),

1437 + )

1438 +

1439 + if idx_mapping_layout == "contiguous":

1440 + idx_mapping = torch.arange(num_reqs, dtype=torch.int32).cuda()

1441 + else:

1442 + idx_mapping = (

1443 + torch.arange(num_reqs, dtype=torch.int32)

1444 + .repeat_interleave(num_steps)

1445 + .view(num_reqs, num_steps)

1446 + .cuda()

1447 + )[:, 1]

1448 + assert idx_mapping.stride(0) == num_steps

1449 +

1450 + skip_masks: list[torch.Tensor] = []

1451 + original = philox_gumbel_sample

1452 +

1453 + def spy(*args, **kwargs):

1454 + skip_masks.append(kwargs["skip_mask"].cpu())

1455 + return original(*args, **kwargs)

1456 +

1457 + monkeypatch.setattr("vllm.v1.worker.gpu.sample.watermark.philox_gumbel_sample", spy)

1458 +

1459 + draft_logits = torch.zeros(num_reqs, num_steps, vocab_size + 1).cuda()

1460 + step_logits = []

1461 + for step in range(num_steps):

1462 + # Token 6 wins by a margin no Gumbel draw can close, on either path.

1463 + logits = torch.zeros(num_reqs, vocab_size).cuda()

1464 + logits[:, 6] = 100.0

1465 + logits[:, 0] = float(step) + 1.0

1466 + step_logits.append(logits)

1467 + sampled = draft_watermarker.sample(

1468 + logits,

1469 + idx_mapping,

1470 + temperature=torch.ones(num_reqs).cuda(),

1471 + seed=torch.zeros(num_reqs, dtype=torch.int64).cuda(),

1472 + pos=torch.full((num_reqs,), 5 + step, dtype=torch.int64).cuda(),

1473 + apply_temperature=True,

1474 + is_drafting=True,

1475 + logits_cache_col=torch.tensor(step).cuda(),

1476 + logits_cache=draft_logits,

1477 + use_fp64=False,

1478 + )

1479 + assert torch.equal(sampled.cpu(), torch.tensor([6, 6]))

1480 +

1481 + # Step 0: request 0's [5, 6] is committed at position 2, request 1's [6, 3]

1482 + # is new. Step 1: [6, 6] and [3, 6] are both new. Step 2: request 0's

1483 + # [6, 6] repeats its own step-1 context and request 1's repeats position 4.

1484 + assert [mask.tolist() for mask in skip_masks] == [

1485 + [True, False],

1486 + [False, False],

1487 + [True, True],

1488 + ]

1489 + assert torch.equal(

1490 + draft_watermarker.prior_contexts[:num_reqs].cpu(),

1491 + torch.tensor([[[5, 6], [6, 6], [6, 6]], [[6, 3], [3, 6], [6, 6]]]),

1492 + )

1493 + assert torch.equal(

1494 + draft_watermarker.contexts[:num_reqs].cpu(), torch.tensor([[6, 6], [6, 6]])

1495 + )

1496 + for step, logits in enumerate(step_logits):

1497 + # The 0-d draft step selects the logits cache column.

1498 + assert torch.equal(draft_logits[:, step, :vocab_size], logits)

1499 + assert not draft_logits[:, :, -1].any()

1500 +

1501 +

1502 def test_dspark_reduced_vocab_draft_sampler_applies_watermarking(monkeypatch):

1503 speculator = object.__new__(DSparkSpeculator)

1504 speculator.draft_logits = torch.zeros(2, 1, 8)

@@ -982,15 +1511,23 @@ def test_dspark_reduced_vocab_draft_sampler_applies_watermarking(monkeypatch):

1511 speculator.acceptance_estimator = None

1512 watermark_logits: list[torch.Tensor] = []

1513

985 - def sample(logits, sampled, idx_map, temperature):

1514 + def sample(

1515 + logits,

1516 + idx_mapping,

1517 + temperature,

1518 + seed,

1519 + pos,

1520 + apply_temperature,

1521 + is_drafting,

1522 + logits_cache=None,

1523 + logits_cache_col=None,

1524 + use_fp64=False,

1525 + logits_cache_source=None,

1526 + ):

1527 watermark_logits.append(logits.clone())

987 - return sampled + 1

1528 + return torch.tensor([4, 5])

1529

1530 speculator.draft_watermarker = SimpleNamespace(sample=sample)

990 - monkeypatch.setattr(

991 - "vllm.v1.worker.gpu.spec_decode.dspark.speculator.gumbel_sample",

992 - lambda *args, **kwargs: torch.tensor([3, 4]),

993 - )

1531

1532 sampled = speculator._sample_logits(

1533 torch.tensor([[10.0, 20.0], [30.0, 40.0]]),

vllm/config/vllm.py

-7

@@ -1344,13 +1344,6 @@ class VllmConfig:

1344 "share of output tokens supplied by accepted drafts.",

1345 scope="global",

1346 )

1347 - if watermark_config.deduplicate_contexts != "none":

1348 - logger.warning_once(

1349 - "Context deduplication is not supported with speculative "

1350 - "decoding and will not be applied to accepted drafts, "

1351 - "rejection-recovery tokens, or bonus tokens.",

1352 - scope="global",

1353 - )

1347 if watermark_config.algorithm == "dual_key_gumbel" and (

1348 watermark_config.alpha != get_field(WatermarkConfig, "alpha").default

1349 ):

vllm/v1/watermarking/gpu_sampler.py

+11 -2

@@ -119,8 +119,9 @@ class GPUWatermarkSampler(Sampler):

119 self,

120 expanded_idx_mapping: torch.Tensor,

121 contexts: torch.Tensor,

122 + expanded_local_pos: torch.Tensor | None = None,

123 ) -> torch.Tensor:

123 - return repeated_context_mask(

124 + repeated = repeated_context_mask(

125 self.req_states.all_token_ids.gpu,

126 expanded_idx_mapping,

127 self.req_states.prompt_len.gpu,

@@ -128,8 +129,16 @@ class GPUWatermarkSampler(Sampler):

129 contexts,

130 self.deduplicate_contexts_max_history,

131 include_prompt=self.deduplicate_contexts == "all",

131 - skip_partial_context=self.deduplicate_contexts == "all",

132 + skip_partial_context=(

133 + self.deduplicate_contexts == "all" and expanded_local_pos is None

134 + ),

135 + history_offsets=expanded_local_pos,

136 + local_positions=expanded_local_pos,

137 + num_speculative_steps=(

138 + self.num_speculative_tokens if expanded_local_pos is not None else 0

139 + ),

140 )

141 + return repeated

142

143 def _get_contexts(

144 self,

vllm/v1/watermarking/gumbel.py

@@ -94,6 +94,10 @@ class GumbelWatermarker(Watermarker):

94 seeds=random_sampler.seeds,

95 positions=random_sampler.positions,

96 use_fp64=random_sampler.use_fp64,

97 + is_drafting=random_sampler.is_drafting,

98 + logits_cache=random_sampler.logits_cache,

99 + logits_cache_col=random_sampler.logits_cache_col,

100 + logits_cache_source=random_sampler.logits_cache_source,

101 )

102 return WatermarkSample(token_ids, logits)

103

vllm/v1/watermarking/spec_decode.py

+112 -10

@@ -3,14 +3,16 @@

4 import torch

6 -from vllm.config.watermarking import WatermarkConfig

6 +from vllm.config.watermarking import WatermarkConfig, WatermarkContextScope

7 from vllm.v1.watermarking.factory import create_watermarker

8 from vllm.v1.watermarking.gumbel import GumbelWatermarker

9 from vllm.v1.watermarking.prfs import PhiloxPRF

10 from vllm.v1.watermarking.watermarker import (

11 + RandomSampler,

12 SupportsSpeculativeDecoding,

13 Watermarker,

14 )

15 +from vllm.v1.worker.gpu.sample.watermark import draft_watermarking_mask

16 from vllm.v1.worker.gpu.spec_decode.rejection_sampler_utils import rejection_sample

@@ -20,8 +22,14 @@ class DraftWatermarker:

22 watermarker: Watermarker,

23 max_num_reqs: int,

24 device: torch.device,

25 + num_speculative_steps: int,

26 + deduplicate_contexts: WatermarkContextScope,

27 + deduplicate_contexts_max_history: int | None,

28 ) -> None:

29 self.watermarker = watermarker

30 + self.num_speculative_steps = num_speculative_steps

31 + self.deduplicate_contexts = deduplicate_contexts

32 + self.deduplicate_contexts_max_history = deduplicate_contexts_max_history

33 self.contexts = torch.zeros(

34 max_num_reqs,

35 watermarker.context_width,

@@ -29,30 +37,112 @@ class DraftWatermarker:

37 device=device,

38 )

39 self.enabled = torch.zeros(max_num_reqs, dtype=torch.bool, device=device)

40 + self.prior_contexts = torch.zeros(

41 + max_num_reqs,

42 + num_speculative_steps,

43 + watermarker.context_width,

44 + dtype=torch.int64,

45 + device=device,

46 + )

47 + self.all_token_ids: torch.Tensor | None = None

48 + self.prompt_lens: torch.Tensor | None = None

49 + self.total_lens: torch.Tensor | None = None

33 - def prepare(self, contexts: torch.Tensor, enabled: torch.Tensor) -> None:

51 + def prepare(

52 + self,

53 + contexts: torch.Tensor,

54 + enabled: torch.Tensor,

55 + all_token_ids: torch.Tensor,

56 + prompt_lens: torch.Tensor,

57 + total_lens: torch.Tensor,

58 + ) -> None:

59 num_reqs = contexts.shape[0]

60 self.contexts[:num_reqs].copy_(contexts)

61 self.enabled[:num_reqs].copy_(enabled)

62 + self.all_token_ids = all_token_ids

63 + self.prompt_lens = prompt_lens

64 + self.total_lens = total_lens

65 +

66 + def _sampling_state(

67 + self,

68 + logits: torch.Tensor,

69 + idx_mapping: torch.Tensor,

70 + request_temperatures: torch.Tensor,

71 + draft_step: int | torch.Tensor,

72 + contexts: torch.Tensor,

73 + ) -> tuple[torch.Tensor, torch.Tensor]:

74 + num_rows = logits.shape[0]

75 + steps = torch.as_tensor(draft_step, device=logits.device, dtype=torch.int64)

76 + if steps.ndim == 0:

77 + steps = steps.expand(num_rows)

78 + enabled = self.enabled[:num_rows] & (request_temperatures != 0)

79 + if self.deduplicate_contexts == "none":

80 + return steps, enabled

81 + assert self.all_token_ids is not None

82 + assert self.prompt_lens is not None

83 + assert self.total_lens is not None

84 + enabled = draft_watermarking_mask(

85 + self.all_token_ids,

86 + idx_mapping,

87 + self.prompt_lens,

88 + self.total_lens,

89 + self.prior_contexts,

90 + contexts,

91 + steps,

92 + enabled,

93 + self.deduplicate_contexts_max_history,

94 + include_prompt=self.deduplicate_contexts == "all",

95 + )

96 + return steps, enabled

98 def sample(

99 self,

100 logits: torch.Tensor,

41 - ordinary_sampled: torch.Tensor,

101 idx_mapping: torch.Tensor,

102 temperature: torch.Tensor,

103 + seed: torch.Tensor,

104 + pos: torch.Tensor,

105 + apply_temperature: bool,

106 + is_drafting: bool,

107 + logits_cache: torch.Tensor | None = None,

108 + logits_cache_col: torch.Tensor | None = None,

109 + use_fp64: bool = False,

110 + logits_cache_source: torch.Tensor | None = None,

111 ) -> torch.Tensor:

112 + assert apply_temperature

113 + assert is_drafting

114 + assert logits_cache is not None

115 + assert logits_cache_col is not None

116 + num_rows = logits.shape[0]

117 request_temperatures = temperature[idx_mapping]

118 + contexts = self.contexts[:num_rows]

119 + steps, enabled = self._sampling_state(

120 + logits, idx_mapping, request_temperatures, logits_cache_col, contexts

121 + )

122 +

123 processed_logits = logits / torch.where(

124 request_temperatures == 0, 1, request_temperatures

125 ).unsqueeze(-1)

49 - watermarked = self.watermarker.sample(

126 + random_sampler = RandomSampler(

127 + expanded_idx_mapping=idx_mapping,

128 + temperatures=temperature,

129 + seeds=seed,

130 + positions=pos,

131 + use_fp64=use_fp64,

132 + is_drafting=True,

133 + logits_cache=logits_cache,

134 + logits_cache_col=steps,

135 + logits_cache_source=(

136 + logits if logits_cache_source is None else logits_cache_source

137 + ),

138 + )

139 + sampled = self.watermarker.sample(

140 processed_logits,

51 - self.contexts[: logits.shape[0]],

141 + contexts,

142 + random_sampler=random_sampler,

143 + skip_mask=~enabled,

144 ).token_ids

53 - enabled = self.enabled[: logits.shape[0]] & (request_temperatures != 0)

54 - sampled = torch.where(enabled, watermarked, ordinary_sampled)

55 - contexts = self.contexts[: logits.shape[0]]

145 +

146 contexts.copy_(torch.cat((contexts[:, 1:], sampled.unsqueeze(-1)), dim=-1))

147 return sampled

148

@@ -68,9 +158,19 @@ def create_speculative_draft_watermarker(

158 max_num_reqs: int,

159 device: torch.device,

160 allow_target_only: bool,

161 + num_speculative_steps: int = 1,

162 + deduplicate_contexts: WatermarkContextScope = "none",

163 + deduplicate_contexts_max_history: int | None = 8192,

164 ) -> DraftWatermarker | None:

165 if isinstance(watermarker, SupportsSpeculativeDecoding):

73 - return DraftWatermarker(watermarker.draft_watermarker, max_num_reqs, device)

166 + return DraftWatermarker(

167 + watermarker.draft_watermarker,

168 + max_num_reqs,

169 + device,

170 + num_speculative_steps,

171 + deduplicate_contexts,

172 + deduplicate_contexts_max_history,

173 + )

174 if allow_target_only:

175 return None

176 raise ValueError(

@@ -130,9 +230,10 @@ def watermarked_rejection_sample(

230 contexts: torch.Tensor,

231 watermarking: torch.Tensor,

232 watermarker: Watermarker,

233 + watermarking_skip_mask: torch.Tensor | None = None,

234 use_fp64: bool = False,

235 ) -> tuple[torch.Tensor, torch.Tensor]:

135 - """Use the target key for rejection recovery and bonus tokens."""

236 + """Use the target key except at contexts selected for ordinary sampling."""

237 assert contexts.shape == (target_logits.shape[0], watermarker.context_width)

238 return rejection_sample(

239 target_logits,

@@ -149,5 +250,6 @@ def watermarked_rejection_sample(

250 use_fp64=use_fp64,

251 contexts=contexts,

252 watermarking=watermarking,

253 + watermarking_skip_mask=watermarking_skip_mask,

254 watermark_key=_resolve_watermark_key(watermarker),

255 )

vllm/v1/watermarking/watermarker.py

+8 -1

@@ -23,6 +23,10 @@ class RandomSampler:

23 seeds: torch.Tensor

24 positions: torch.Tensor

25 use_fp64: bool = False

26 + is_drafting: bool = False

27 + logits_cache: torch.Tensor | None = None

28 + logits_cache_col: torch.Tensor | None = None

29 + logits_cache_source: torch.Tensor | None = None

31 def __call__(self, logits: torch.Tensor) -> torch.Tensor:

32 return gumbel_sample(

@@ -32,7 +36,10 @@ class RandomSampler:

36 self.seeds,

37 self.positions,

38 apply_temperature=False,

35 - is_drafting=False,

39 + is_drafting=self.is_drafting,

40 + logits_cache=self.logits_cache,

41 + logits_cache_col=self.logits_cache_col,

42 + logits_cache_source=self.logits_cache_source,

43 use_fp64=self.use_fp64,

44 )

vllm/v1/worker/gpu/model_runner.py

+2 -4

@@ -914,8 +914,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

914 spec_hidden_states = pre_hc_hidden_states[: hidden_states.shape[0]] # type: ignore[union-attr]

915 if isinstance(self.sampler, GPUWatermarkSampler):

916 self.speculator.prepare_watermarking(

917 - self.sampler._get_contexts(input_batch.idx_mapping),

918 - self.sampler.watermarking.gpu[input_batch.idx_mapping],

917 + self.sampler, input_batch.idx_mapping

918 )

919 with use_workspace_lane(self._draft_workspace_lane):

920 self.speculator.propose(

@@ -2194,8 +2193,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

2193 spec_hidden_states = pre_hc_hidden_states[: draft_hidden_states.size(0)]

2194 if isinstance(self.sampler, GPUWatermarkSampler):

2195 self.speculator.prepare_watermarking(

2197 - self.sampler._get_contexts(input_batch.idx_mapping),

2198 - self.sampler.watermarking.gpu[input_batch.idx_mapping],

2196 + self.sampler, input_batch.idx_mapping

2197 )

2198 with use_workspace_lane(self._draft_workspace_lane):

2199 draft_tokens = self.speculator.propose(

vllm/v1/worker/gpu/sample/gumbel.py

+29 -1

@@ -220,6 +220,8 @@ def gumbel_block_argmax(

220 logits_cache_stride_0,

221 logits_cache_stride_1,

222 logits_cache_col_ptr,

223 + logits_cache_source_ptr,

224 + logits_cache_source_stride,

225 vocab_size,

226 IS_DRAFTING: tl.constexpr,

227 APPLY_TEMPERATURE: tl.constexpr,

@@ -240,12 +242,16 @@ def gumbel_block_argmax(

242 col = tl.load(logits_cache_col_ptr + token_idx)

243 else:

244 col = tl.load(logits_cache_col_ptr)

245 + cached_logits = tl.load(

246 + logits_cache_source_ptr + token_idx * logits_cache_source_stride + block,

247 + mask=mask,

248 + )

249 tl.store(

250 logits_cache_ptr

251 + req_state_idx * logits_cache_stride_0

252 + col * logits_cache_stride_1

253 + block,

248 - logits,

254 + cached_logits,

255 mask=mask & is_valid_req,

256 )

257

@@ -275,6 +281,8 @@ def _gumbel_sample_kernel(

281 logits_cache_stride_0,

282 logits_cache_stride_1,

283 logits_cache_col_ptr,

284 + logits_cache_source_ptr,

285 + logits_cache_source_stride,

286 logits_ptr,

287 logits_stride,

288 expanded_idx_mapping_ptr,

@@ -312,6 +320,8 @@ def _gumbel_sample_kernel(

320 logits_cache_stride_0,

321 logits_cache_stride_1,

322 logits_cache_col_ptr,

323 + logits_cache_source_ptr,

324 + logits_cache_source_stride,

325 vocab_size,

326 IS_DRAFTING=IS_DRAFTING,

327 APPLY_TEMPERATURE=APPLY_TEMPERATURE,

@@ -334,6 +344,7 @@ def gumbel_sample(

344 logits_cache: torch.Tensor | None = None, # [max_num_reqs, num_cols, vocab_size]

345 logits_cache_col: torch.Tensor | None = None, # scalar or [num_tokens]

346 use_fp64: bool = False,

347 + logits_cache_source: torch.Tensor | None = None,

348 ) -> torch.Tensor:

349 # Enforce contiguity on non-strided input tensors

350 expanded_idx_mapping = expanded_idx_mapping.contiguous()

@@ -342,11 +353,26 @@ def gumbel_sample(

353 logits_cache_col = logits_cache_col.contiguous()

354 num_tokens, vocab_size = logits.shape

355 if logits_cache is not None:

356 + if logits_cache_source is None:

357 + logits_cache_source = logits

358 + assert logits_cache_source.shape == logits.shape, (

359 + "logits cache source must match sampled logits shape"

360 + )

361 + assert logits_cache_source.device == logits.device, (

362 + "logits cache source must be on the sampled logits device"

363 + )

364 + assert logits_cache_source.dtype == logits_cache.dtype, (

365 + "logits cache source and destination must have the same dtype"

366 + )

367 assert logits_cache.size(-1) >= vocab_size, (

368 f"draft logits cache vocab dim ({logits_cache.size(-1)}) is narrower "

369 f"than the sampled logits ({vocab_size}). Cached logits would be "

370 "truncated."

371 )

372 + elif logits_cache_source is not None:

373 + raise ValueError("logits_cache_source requires logits_cache")

374 + if logits_cache_source is not None and logits_cache_source.stride(-1) != 1:

375 + logits_cache_source = logits_cache_source.contiguous()

376 BLOCK_SIZE = 1024

377 num_blocks = triton.cdiv(vocab_size, BLOCK_SIZE)

378 local_argmax = logits.new_empty(num_tokens, num_blocks, dtype=torch.int64)

@@ -362,6 +388,8 @@ def gumbel_sample(

388 logits_cache.stride(0) if logits_cache is not None else 0,

389 logits_cache.stride(1) if logits_cache is not None else 0,

390 logits_cache_col,

391 + logits_cache_source,

392 + logits_cache_source.stride(0) if logits_cache_source is not None else 0,

393 logits,

394 logits.stride(0),

395 expanded_idx_mapping,

vllm/v1/worker/gpu/sample/watermark.py

+315 -7

@@ -51,12 +51,25 @@ def _repeated_context_mask_kernel(

51 req_indices_ptr,

52 prompt_lens_ptr,

53 total_lens_ptr,

54 + history_offsets_ptr,

55 + history_offsets_stride,

56 contexts_ptr,

57 context_stride,

58 + local_positions_ptr,

59 + prior_contexts_ptr,

60 + prior_contexts_stride_0,

61 + prior_contexts_stride_1,

62 + steps_ptr,

63 + steps_stride,

64 + enabled_ptr,

65 CONTEXT_WIDTH: tl.constexpr,

66 + NUM_SPECULATIVE_STEPS: tl.constexpr,

67 MAX_HISTORY: tl.constexpr,

68 INCLUDE_PROMPT: tl.constexpr,

69 SKIP_PARTIAL_CONTEXT: tl.constexpr,

70 + HAS_HISTORY_OFFSETS: tl.constexpr,

71 + SPECULATIVE_CONTEXTS: tl.constexpr,

72 + DRAFT_CONTEXTS: tl.constexpr,

73 BLOCK: tl.constexpr,

74 ):

75 row = tl.program_id(0).to(tl.int64)

@@ -71,12 +84,16 @@ def _repeated_context_mask_kernel(

84 return

85 sequence_start = 0 if INCLUDE_PROMPT else prompt_len

86 history_len = total_len - sequence_start

87 + history_offset = 0

88 + if HAS_HISTORY_OFFSETS:

89 + history_offset = tl.load(history_offsets_ptr + row * history_offsets_stride)

91 offsets = tl.arange(0, BLOCK)

92 repeated = tl.full((), 0, tl.int32)

93 scan_start = 0

94 if MAX_HISTORY > 0:

79 - scan_start = tl.maximum(history_len - MAX_HISTORY, 0)

95 + remaining_history = tl.maximum(MAX_HISTORY - history_offset, 0)

96 + scan_start = tl.maximum(history_len - remaining_history, 0)

97 # Aligned loop bounds let the per-lane offsets fold into immediates.

98 aligned_start = (scan_start // BLOCK) * BLOCK

99 for block_start in tl.range(aligned_start, history_len, BLOCK):

@@ -96,6 +113,93 @@ def _repeated_context_mask_kernel(

113 )

114 matches &= historical_token == context_token

115 repeated |= tl.max(matches.to(tl.int32), axis=0)

116 +

117 + if SPECULATIVE_CONTEXTS:

118 + local_pos = tl.load(local_positions_ptr + row)

119 + if INCLUDE_PROMPT:

120 + partial_context = tl.full((), 0, tl.int32)

121 + for context_offset in range(CONTEXT_WIDTH):

122 + context_token = tl.load(

123 + contexts_ptr + row * context_stride + context_offset

124 + )

125 + partial_context |= context_token < 0

126 + repeated |= partial_context

127 + for offset in range(1, NUM_SPECULATIVE_STEPS + 1):

128 + prior_row = row - offset

129 + in_history = (prior_row >= 0) & (local_pos >= offset)

130 + if MAX_HISTORY > 0:

131 + in_history &= offset <= MAX_HISTORY

132 + prior_req_idx = tl.load(

133 + req_indices_ptr + tl.maximum(prior_row, 0),

134 + mask=in_history,

135 + other=-1,

136 + )

137 + prior_local_pos = tl.load(

138 + local_positions_ptr + tl.maximum(prior_row, 0),

139 + mask=in_history,

140 + other=-1,

141 + )

142 + matches = (

143 + in_history

144 + & (req_idx == prior_req_idx)

145 + & (local_pos == prior_local_pos + offset)

146 + )

147 + for context_offset in range(CONTEXT_WIDTH):

148 + context_token = tl.load(

149 + contexts_ptr + row * context_stride + context_offset

150 + )

151 + prior_token = tl.load(

152 + contexts_ptr

153 + + tl.maximum(prior_row, 0) * context_stride

154 + + context_offset,

155 + mask=in_history,

156 + other=-2,

157 + )

158 + matches &= context_token == prior_token

159 + repeated |= matches

160 +

161 + if DRAFT_CONTEXTS:

162 + step = tl.load(steps_ptr + row * steps_stride).to(tl.int64)

163 + partial_context = tl.full((), 0, tl.int32)

164 + for context_offset in range(CONTEXT_WIDTH):

165 + context_token = tl.load(

166 + contexts_ptr + row * context_stride + context_offset

167 + )

168 + partial_context |= context_token < 0

169 + if INCLUDE_PROMPT:

170 + repeated |= partial_context

171 +

172 + for prior_step in range(NUM_SPECULATIVE_STEPS):

173 + in_history = prior_step < step

174 + if MAX_HISTORY > 0:

175 + in_history &= prior_step >= step - MAX_HISTORY

176 + matches = in_history

177 + for context_offset in range(CONTEXT_WIDTH):

178 + context_token = tl.load(

179 + contexts_ptr + row * context_stride + context_offset

180 + )

181 + prior_token = tl.load(

182 + prior_contexts_ptr

183 + + row * prior_contexts_stride_0

184 + + prior_step * prior_contexts_stride_1

185 + + context_offset

186 + )

187 + matches &= context_token == prior_token

188 + repeated |= matches

189 +

190 + for context_offset in range(CONTEXT_WIDTH):

191 + context_token = tl.load(

192 + contexts_ptr + row * context_stride + context_offset

193 + )

194 + tl.store(

195 + prior_contexts_ptr

196 + + row * prior_contexts_stride_0

197 + + step * prior_contexts_stride_1

198 + + context_offset,

199 + context_token,

200 + )

201 + enabled = tl.load(enabled_ptr + row)

202 + repeated = enabled & (repeated == 0)

203 tl.store(output_ptr + row, repeated)

204

205

@@ -108,6 +212,7 @@ def _repeated_context_mask_cpu(

212 max_history: int | None = None,

213 include_prompt: bool = False,

214 skip_partial_context: bool = False,

215 + history_offsets: torch.Tensor | None = None,

216 ) -> torch.Tensor:

217 """Reference implementation; the parity tests check the Triton kernel against it."""

218 repeated = torch.zeros(len(req_indices), dtype=torch.bool)

@@ -124,8 +229,11 @@ def _repeated_context_mask_cpu(

229 history_tokens = all_token_ids[req_idx, sequence_start:total_len].tolist()

230 prefix = [-1] * contexts.shape[-1]

231 current_context = tuple(contexts[row].tolist())

232 + history_offset = 0 if history_offsets is None else int(history_offsets[row])

233 history_start = (

128 - max(0, len(history_tokens) - max_history) if max_history is not None else 0

234 + max(0, len(history_tokens) - max(0, max_history - history_offset))

235 + if max_history is not None

236 + else 0

237 )

238 for history_pos, token_id in enumerate(history_tokens):

239 if history_pos >= history_start and (

@@ -146,6 +254,9 @@ def repeated_context_mask(

254 max_history: int | None = None,

255 include_prompt: bool = False,

256 skip_partial_context: bool = False,

257 + history_offsets: torch.Tensor | None = None,

258 + local_positions: torch.Tensor | None = None,

259 + num_speculative_steps: int = 0,

260 ) -> torch.Tensor:

261 """Return, per row, whether the row's context already occurred in its history.

262

@@ -163,12 +274,23 @@ def repeated_context_mask(

274 include_prompt: Search the prompt as well as the generated tokens.

275 skip_partial_context: Mark contexts containing start padding so they use

276 ordinary sampling.

277 + history_offsets: Number of newer, non-committed positions preceding each

278 + context. These positions count toward `max_history`. Must be 1-D

279 + over rows; a stride-0 broadcast row is fine.

280 + local_positions: Position within each request's speculative block. When

281 + provided, compare each context with earlier contexts in that block.

282 + num_speculative_steps: Maximum number of draft tokens in the block.

283

284 """

285 if max_history is not None and max_history < 1:

286 raise ValueError("max_history must be positive or None")

287 + if (local_positions is None) != (num_speculative_steps == 0):

288 + raise ValueError(

289 + "local_positions and positive num_speculative_steps must be "

290 + "provided together"

291 + )

292 if all_token_ids.device.type == "cpu":

171 - return _repeated_context_mask_cpu(

293 + repeated = _repeated_context_mask_cpu(

294 all_token_ids,

295 req_indices,

296 prompt_lens,

@@ -177,10 +299,42 @@ def repeated_context_mask(

299 max_history,

300 include_prompt,

301 skip_partial_context,

302 + history_offsets,

303 )

304 + if local_positions is not None:

305 + max_offset = num_speculative_steps + 1

306 + if max_history is not None:

307 + max_offset = min(max_offset, max_history + 1)

308 + for offset in range(1, max_offset):

309 + prior_contexts = torch.cat(

310 + (contexts[:offset], contexts[:-offset]), dim=0

311 + )

312 + prior_requests = torch.cat(

313 + (req_indices[:offset], req_indices[:-offset]), dim=0

314 + )

315 + prior_local_pos = torch.cat(

316 + (local_positions[:offset], local_positions[:-offset]), dim=0

317 + )

318 + repeated |= (

319 + (req_indices >= 0)

320 + & (local_positions >= offset)

321 + & (req_indices == prior_requests)

322 + & (local_positions == prior_local_pos + offset)

323 + & (contexts == prior_contexts).all(dim=-1)

324 + )

325 + if include_prompt:

326 + repeated |= (contexts < 0).any(dim=-1)

327 + return repeated

328

329 if contexts.stride(-1) != 1:

330 contexts = contexts.contiguous()

331 + # Only history_offsets passes its row stride to the kernel.

332 + req_indices = req_indices.contiguous()

333 + if local_positions is not None:

334 + local_positions = local_positions.contiguous()

335 + # `prompt_lens` and `total_lens` are indexed flat by request slot, and

336 + # `all_token_ids` takes a row stride but is read flat within the row. Those

337 + # three are contiguous by contract; every caller owns them outright.

338 repeated = torch.empty(len(req_indices), dtype=torch.bool, device=contexts.device)

339 _repeated_context_mask_kernel[(len(req_indices),)](

340 repeated,

@@ -189,17 +343,104 @@ def repeated_context_mask(

343 req_indices,

344 prompt_lens,

345 total_lens,

346 + history_offsets,

347 + 0 if history_offsets is None else history_offsets.stride(0),

348 contexts,

349 contexts.stride(0),

350 + local_positions,

351 + None,

352 + 0,

353 + 0,

354 + None,

355 + 0,

356 + None,

357 CONTEXT_WIDTH=contexts.shape[-1],

358 + NUM_SPECULATIVE_STEPS=num_speculative_steps,

359 MAX_HISTORY=0 if max_history is None else max_history,

360 INCLUDE_PROMPT=include_prompt,

361 SKIP_PARTIAL_CONTEXT=skip_partial_context,

362 + HAS_HISTORY_OFFSETS=history_offsets is not None,

363 + SPECULATIVE_CONTEXTS=local_positions is not None,

364 + DRAFT_CONTEXTS=False,

365 BLOCK=512,

366 )

367 return repeated

368

369

370 +def draft_watermarking_mask(

371 + all_token_ids: torch.Tensor,

372 + req_indices: torch.Tensor,

373 + prompt_lens: torch.Tensor,

374 + total_lens: torch.Tensor,

375 + prior_contexts: torch.Tensor,

376 + contexts: torch.Tensor,

377 + steps: torch.Tensor,

378 + enabled: torch.Tensor,

379 + max_history: int | None,

380 + include_prompt: bool,

381 +) -> torch.Tensor:

382 + if contexts.device.type == "cpu":

383 + repeated = repeated_context_mask(

384 + all_token_ids,

385 + req_indices,

386 + prompt_lens,

387 + total_lens,

388 + contexts,

389 + max_history,

390 + include_prompt=include_prompt,

391 + history_offsets=steps,

392 + )

393 + prior_steps = torch.arange(prior_contexts.shape[1]).unsqueeze(0)

394 + in_history = prior_steps < steps.unsqueeze(1)

395 + if max_history is not None:

396 + in_history &= prior_steps >= steps.unsqueeze(1) - max_history

397 + repeated |= (

398 + (prior_contexts[: len(contexts)] == contexts.unsqueeze(1)).all(dim=-1)

399 + & in_history

400 + ).any(dim=-1)

401 + if include_prompt:

402 + repeated |= (contexts < 0).any(dim=-1)

403 + row_indices = torch.arange(len(contexts))

404 + prior_contexts.index_put_((row_indices, steps), contexts)

405 + return enabled & ~repeated

406 +

407 + if contexts.stride(-1) != 1:

408 + contexts = contexts.contiguous()

409 + # Only `steps` passes its row stride to the kernel.

410 + req_indices = req_indices.contiguous()

411 + enabled = enabled.contiguous()

412 + active = torch.empty(len(contexts), dtype=torch.bool, device=contexts.device)

413 + _repeated_context_mask_kernel[(len(contexts),)](

414 + active,

415 + all_token_ids,

416 + all_token_ids.stride(0),

417 + req_indices,

418 + prompt_lens,

419 + total_lens,

420 + steps,

421 + steps.stride(0),

422 + contexts,

423 + contexts.stride(0),

424 + None,

425 + prior_contexts,

426 + prior_contexts.stride(0),

427 + prior_contexts.stride(1),

428 + steps,

429 + steps.stride(0),

430 + enabled,

431 + NUM_SPECULATIVE_STEPS=prior_contexts.shape[1],

432 + CONTEXT_WIDTH=contexts.shape[1],

433 + MAX_HISTORY=0 if max_history is None else max_history,

434 + INCLUDE_PROMPT=include_prompt,

435 + SKIP_PARTIAL_CONTEXT=False,

436 + HAS_HISTORY_OFFSETS=True,

437 + SPECULATIVE_CONTEXTS=False,

438 + DRAFT_CONTEXTS=True,

439 + BLOCK=512,

440 + )

441 + return active

442 +

443 +

444 @triton.jit

445 def _mulhilo32(multiplier: tl.constexpr, value):

446 high = tl_math.umulhi(multiplier, value)

@@ -351,28 +592,58 @@ def _philox_gumbel_kernel(

592 seeds_ptr,

593 pos_ptr,

594 temp_ptr,

595 + logits_cache_ptr,

596 + logits_cache_stride_0,

597 + logits_cache_stride_1,

598 + logits_cache_col_ptr,

599 + logits_cache_source_ptr,

600 + logits_cache_source_stride,

601 key_0_value,

602 key_1_value,

603 vocab_size,

604 CONTEXT_WIDTH: tl.constexpr,

605 BLOCK_SIZE: tl.constexpr,

606 USE_FP64: tl.constexpr,

607 + IS_DRAFTING: tl.constexpr,

608 + PER_TOKEN_COL: tl.constexpr,

609 ):

610 row = tl.program_id(0).to(tl.int64)

611 block_index = tl.program_id(1)

612 groups = block_index * (BLOCK_SIZE // 4) + tl.arange(0, BLOCK_SIZE // 4)

613 + candidate = block_index * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)

614 + mask = candidate < vocab_size

615

365 - if skip_mask_ptr is not None:

616 + req_state_idx = row

617 + valid_req = True

618 + if expanded_idx_mapping_ptr is not None:

619 req_state_idx = tl.load(expanded_idx_mapping_ptr + row).to(tl.int64)

620 valid_req = req_state_idx >= 0

621 +

622 + if logits_cache_ptr is not None:

623 + if PER_TOKEN_COL:

624 + col = tl.load(logits_cache_col_ptr + row)

625 + else:

626 + col = tl.load(logits_cache_col_ptr)

627 + cached_logits = tl.load(

628 + logits_cache_source_ptr + row * logits_cache_source_stride + candidate,

629 + mask=mask,

630 + )

631 + tl.store(

632 + logits_cache_ptr

633 + + req_state_idx * logits_cache_stride_0

634 + + col * logits_cache_stride_1

635 + + candidate,

636 + cached_logits,

637 + mask=mask & valid_req,

638 + )

639 +

640 + if skip_mask_ptr is not None: # noqa: SIM102

641 if tl.load(skip_mask_ptr + row):

642 temp = tl.load(temp_ptr + req_state_idx, mask=valid_req, other=0.0).to(

643 tl.float32

644 )

645 seed = tl.load(seeds_ptr + req_state_idx, mask=valid_req, other=0)

646 pos = tl.load(pos_ptr + row)

374 - candidate = block_index * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)

375 - mask = candidate < vocab_size

647 logits_row = logits_ptr + row * logits_stride

648 logits = tl.load(logits_row + candidate, mask=mask, other=float("-inf")).to(

649 tl.float32

@@ -384,7 +655,7 @@ def _philox_gumbel_kernel(

655 seed,

656 pos,

657 temp,

387 - IS_DRAFTING=False,

658 + IS_DRAFTING=IS_DRAFTING,

659 USE_FP64=USE_FP64,

660 APPLY_TEMPERATURE=False,

661 )

@@ -443,12 +714,37 @@ def philox_gumbel_sample(

714 seeds: torch.Tensor | None = None,

715 positions: torch.Tensor | None = None,

716 use_fp64: bool = False,

717 + is_drafting: bool = False,

718 + logits_cache: torch.Tensor | None = None,

719 + logits_cache_col: torch.Tensor | None = None,

720 + logits_cache_source: torch.Tensor | None = None,

721 ) -> torch.Tensor:

722 sampling_state = (expanded_idx_mapping, temperatures, seeds, positions)

723 if skip_mask is None and any(value is not None for value in sampling_state):

724 raise ValueError("sampling state requires skip_mask")

725 if skip_mask is not None and any(value is None for value in sampling_state):

726 raise ValueError("skip_mask requires complete sampling state")

727 + if logits_cache is not None:

728 + if logits_cache_col is None or logits_cache_source is None:

729 + raise ValueError("logits cache requires its column and source logits")

730 + if expanded_idx_mapping is None:

731 + raise ValueError("logits cache requires expanded_idx_mapping")

732 + assert logits_cache_source.shape == logits.shape, (

733 + "logits cache source must match sampled logits shape"

734 + )

735 + assert logits_cache_source.device == logits.device, (

736 + "logits cache source must be on the sampled logits device"

737 + )

738 + assert logits_cache_source.dtype == logits_cache.dtype, (

739 + "logits cache source and destination must have the same dtype"

740 + )

741 + assert logits_cache.size(-1) >= logits.shape[-1], (

742 + f"draft logits cache vocab dim ({logits_cache.size(-1)}) is narrower "

743 + f"than the sampled logits ({logits.shape[-1]}). Cached logits would be "

744 + "truncated."

745 + )

746 + elif logits_cache_col is not None or logits_cache_source is not None:

747 + raise ValueError("logits cache column and source require logits_cache")

748

749 if logits.stride(-1) != 1:

750 logits = logits.contiguous()

@@ -460,6 +756,10 @@ def philox_gumbel_sample(

756 expanded_idx_mapping = expanded_idx_mapping.contiguous()

757 if positions is not None:

758 positions = positions.contiguous()

759 + if logits_cache_col is not None:

760 + logits_cache_col = logits_cache_col.contiguous()

761 + if logits_cache_source is not None and logits_cache_source.stride(-1) != 1:

762 + logits_cache_source = logits_cache_source.contiguous()

763 num_tokens, vocab_size = logits.shape

764 block_size = 1024

765 num_blocks = triton.cdiv(vocab_size, block_size)

@@ -483,12 +783,20 @@ def philox_gumbel_sample(

783 seeds,

784 positions,

785 temperatures,

786 + logits_cache,

787 + logits_cache.stride(0) if logits_cache is not None else 0,

788 + logits_cache.stride(1) if logits_cache is not None else 0,

789 + logits_cache_col,

790 + logits_cache_source,

791 + logits_cache_source.stride(0) if logits_cache_source is not None else 0,

792 key & _UINT32_MASK_VALUE,

793 key >> 32,

794 vocab_size,

795 CONTEXT_WIDTH=contexts.shape[-1],

796 BLOCK_SIZE=block_size,

797 USE_FP64=use_fp64,

798 + IS_DRAFTING=is_drafting,

799 + PER_TOKEN_COL=logits_cache_col is not None and logits_cache_col.dim() > 0,

800 )

801 max_block_index = local_max.argmax(dim=-1, keepdim=True)

802 return local_argmax.gather(dim=-1, index=max_block_index).view(-1)

vllm/v1/worker/gpu/spec_decode/dspark/speculator.py

+6 -6

@@ -135,7 +135,12 @@ class DSparkSpeculator(DFlashSpeculator):

135

136 # sample_pos is the predicted token's position P. Sampling keys a draw

137 # by the position before the sampled token, P-1.

138 - sampled = gumbel_sample(

138 + sampler = (

139 + gumbel_sample

140 + if self.draft_watermarker is None

141 + else self.draft_watermarker.sample

142 + )

143 + return sampler(

144 logits,

145 idx_map,

146 self.temperature,

@@ -147,11 +152,6 @@ class DSparkSpeculator(DFlashSpeculator):

152 logits_cache_col=self._step_cols[step],

153 use_fp64=self.use_fp64_gumbel,

154 )

150 - if self.draft_watermarker is not None:

151 - sampled = self.draft_watermarker.sample(

152 - logits, sampled, idx_map, self.temperature

153 - )

154 - return sampled

155

156 def _sample_sequential(self, num_reqs: int, head_hidden: torch.Tensor) -> None:

157 if self._draft_topk is not None:

vllm/v1/worker/gpu/spec_decode/rejection_sampler.py

+10 -3

@@ -226,11 +226,18 @@ class RejectionSampler:

226 if self.watermark_key is None:

227 return {}

228 assert isinstance(self.sampler, GPUWatermarkSampler)

229 + contexts = self.sampler._get_contexts(

230 + expanded_idx_mapping, expanded_local_pos, draft_sampled

231 + )

232 + watermarking_skip_mask = None

233 + if self.sampler.deduplicate_contexts != "none":

234 + watermarking_skip_mask = self.sampler._get_repeated_contexts(

235 + expanded_idx_mapping, contexts, expanded_local_pos

236 + )

237 return {

230 - "contexts": self.sampler._get_contexts(

231 - expanded_idx_mapping, expanded_local_pos, draft_sampled

232 - ),

238 + "contexts": contexts,

239 "watermarking": self.sampler.watermarking.gpu,

240 + "watermarking_skip_mask": watermarking_skip_mask,

241 "watermark_key": self.watermark_key,

242 }

243

vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py

+37 -10

@@ -564,7 +564,6 @@ def _rejection_kernel(

564 # stores the target argmax upon first rejection, so it rejects the

565 # placeholder via `accepted` instead.

566 verifying &= is_valid_draft

567 -

567 if verifying:

568 pos = tl.load(pos_ptr + logit_idx)

569 u = tl_rand32(seed, pos, includes_zero=False)

@@ -719,6 +718,8 @@ def _seeded_resample_argmax(

718 0, # logits_cache_stride_0

719 0, # logits_cache_stride_1

720 None, # logits_cache_col_ptr

721 + None, # logits_cache_source_ptr

722 + 0, # logits_cache_source_stride

723 vocab_size,

724 IS_DRAFTING=False,

725 APPLY_TEMPERATURE=False,

@@ -766,6 +767,7 @@ def _resample_kernel(

767 contexts_stride,

768 # [max_num_reqs], uint8 view of a bool tensor

769 watermarking_ptr,

770 + watermarking_skip_mask_ptr,

771 watermark_key_0,

772 watermark_key_1,

773 vocab_size,

@@ -775,6 +777,7 @@ def _resample_kernel(

777 USE_BLOCK_VERIFICATION: tl.constexpr,

778 CONTEXT_WIDTH: tl.constexpr,

779 WATERMARK: tl.constexpr,

780 + DEDUPLICATE_CONTEXTS: tl.constexpr,

781 ):

782 req_idx = tl.program_id(0)

783 resample_idx = tl.load(rejected_step_ptr + req_idx)

@@ -807,6 +810,25 @@ def _resample_kernel(

810 ).to(tl.float32)

811

812 # Compute the residual logits to resample the rejected token from.

813 + is_watermarked = False

814 + if WATERMARK:

815 + is_watermarked = (

816 + tl.load(

817 + watermarking_ptr + req_state_idx,

818 + mask=req_state_idx >= 0,

819 + other=0,

820 + )

821 + != 0

822 + )

823 +

824 + skip_watermarking = False

825 + if DEDUPLICATE_CONTEXTS:

826 + skip_watermarking = (

827 + (tl.load(watermarking_skip_mask_ptr + resample_token_idx) != 0)

828 + & is_watermarked

829 + & (temp != 0.0)

830 + )

831 +

832 if is_bonus or not is_valid_rejected_draft:

833 # Bonus token (no rejections) or -1 placeholder token. In either case,

834 # directly use the target logits.

@@ -868,15 +890,7 @@ def _resample_kernel(

890 # Resample the rejected/bonus token.

891 if WATERMARK:

892 # Padded and greedy rows retain the stock draw.

871 - is_watermarked = (

872 - tl.load(

873 - watermarking_ptr + req_state_idx,

874 - mask=req_state_idx >= 0,

875 - other=0,

876 - )

877 - != 0

878 - )

879 - if is_watermarked & (temp != 0.0):

893 + if is_watermarked & (temp != 0.0) & ~skip_watermarking:

894 watermark_value, idx = philox_gumbel_block_argmax(

895 residual_logits,

896 mask,

@@ -1023,6 +1037,7 @@ def rejection_sample(

1037 use_block_verification: bool = False,

1038 contexts: torch.Tensor | None = None,

1039 watermarking: torch.Tensor | None = None,

1040 + watermarking_skip_mask: torch.Tensor | None = None,

1041 watermark_key: int | None = None,

1042 ) -> tuple[torch.Tensor, torch.Tensor]:

1043 assert target_logits.ndim == 2 and target_logits.stride(-1) == 1

@@ -1036,9 +1051,13 @@ def rejection_sample(

1051 assert watermark == (watermarking is not None) == (watermark_key is not None), (

1052 "contexts, watermarking and watermark_key must be set together."

1053 )

1054 + assert watermarking_skip_mask is None or watermark, (

1055 + "watermarking_skip_mask requires watermarking."

1056 + )

1057 contexts_stride = 0

1058 context_width = 1

1059 watermarking_bytes: torch.Tensor | None = None

1060 + watermarking_skip_mask_bytes: torch.Tensor | None = None

1061 watermark_key_0 = 0

1062 watermark_key_1 = 0

1063 if contexts is not None:

@@ -1053,6 +1072,12 @@ def rejection_sample(

1072 contexts_stride = contexts.stride(0)

1073 context_width = contexts.shape[-1]

1074 watermarking_bytes = watermarking.view(torch.uint8)

1075 + if watermarking_skip_mask is not None:

1076 + assert watermarking_skip_mask.shape == (num_logits,)

1077 + assert watermarking_skip_mask.dtype == torch.bool

1078 + if not watermarking_skip_mask.is_contiguous():

1079 + watermarking_skip_mask = watermarking_skip_mask.contiguous()

1080 + watermarking_skip_mask_bytes = watermarking_skip_mask.view(torch.uint8)

1081 watermark_key_0 = watermark_key & 0xFFFFFFFF

1082 watermark_key_1 = watermark_key >> 32

1083 draft_logits_stride_0 = 0

@@ -1269,6 +1294,7 @@ def rejection_sample(

1294 contexts,

1295 contexts_stride,

1296 watermarking_bytes,

1297 + watermarking_skip_mask_bytes,

1298 watermark_key_0,

1299 watermark_key_1,

1300 vocab_size,

@@ -1278,6 +1304,7 @@ def rejection_sample(

1304 USE_BLOCK_VERIFICATION=use_block_verification,

1305 CONTEXT_WIDTH=context_width,

1306 WATERMARK=watermark,

1307 + DEDUPLICATE_CONTEXTS=watermarking_skip_mask is not None,

1308 )

1309

1310 # Insert the resampled tokens into the output sampled.

vllm/v1/worker/gpu/spec_decode/speculator.py

+20 -7

@@ -38,6 +38,7 @@ from vllm.v1.worker.gpu.spec_decode.acceptance_estimator import (

38 from vllm.v1.worker.utils import AttentionGroup

40 if TYPE_CHECKING:

41 + from vllm.v1.watermarking.gpu_sampler import GPUWatermarkSampler

42 from vllm.v1.worker.gpu.pcp_manager import PCPManager

44 logger = init_logger(__name__)

@@ -198,6 +199,9 @@ class DraftModelSpeculator(BaseSpeculator):

199 self.max_num_reqs,

200 device,

201 watermark_config.allow_target_only_watermarking,

202 + self.num_speculative_steps,

203 + watermark_config.deduplicate_contexts,

204 + watermark_config.deduplicate_contexts_max_history,

205 )

206

207 self.supports_mm_inputs = False

@@ -411,7 +415,12 @@ class DraftModelSpeculator(BaseSpeculator):

415 ) -> torch.Tensor:

416 if draft_logits is not None:

417 logits = self.model.compute_logits(hidden_states)

414 - sampled = gumbel_sample(

418 + sampler = (

419 + gumbel_sample

420 + if self.draft_watermarker is None

421 + else self.draft_watermarker.sample

422 + )

423 + sampled = sampler(

424 logits,

425 idx_mapping,

426 temperature,

@@ -423,10 +432,6 @@ class DraftModelSpeculator(BaseSpeculator):

432 logits_cache_col=draft_step,

433 use_fp64=self.use_fp64_gumbel,

434 )

426 - if self.draft_watermarker is not None:

427 - sampled = self.draft_watermarker.sample(

428 - logits, sampled, idx_mapping, temperature

429 - )

435 elif self.use_local_argmax_reduction:

436 return self.model.get_top_tokens(hidden_states)

437 else:

@@ -462,11 +467,19 @@ class DraftModelSpeculator(BaseSpeculator):

467 self.acceptance_estimator.step(idx_mapping, num_sampled, num_rejected)

468

469 def prepare_watermarking(

465 - self, contexts: torch.Tensor, watermarking: torch.Tensor

470 + self,

471 + sampler: "GPUWatermarkSampler",

472 + idx_mapping: torch.Tensor,

473 ) -> None:

474 if self.draft_watermarker is None:

475 return

469 - self.draft_watermarker.prepare(contexts, watermarking)

476 + self.draft_watermarker.prepare(

477 + sampler._get_contexts(idx_mapping),

478 + sampler.watermarking.gpu[idx_mapping],

479 + sampler.req_states.all_token_ids.gpu,

480 + sampler.req_states.prompt_len.gpu,

481 + sampler.req_states.total_len.gpu,

482 + )

483

484 def _copy_request_inputs(

485 self,