From 202d497aa897edef87fc291b7f428d7dd473aa06 Mon Sep 17 00:00:00 2001 From: simon-veitner-redhat Date: Fri, 2 Oct 2026 20:32:11 +0200 Subject: [PATCH] [Bugfix][Watermarking] Keep draft prompt lengths valid under CUDA graphs (#59779) Signed-off-by: Simon Veitner Co-authored-by: Lucas Wilkinson --- tests/watermarking/test_watermarking.py | 48 +++++++++++++++++++++++++ vllm/v1/watermarking/spec_decode.py | 5 ++- 2 files changed, 50 insertions(+), 3 deletions(-) diff --git a/tests/watermarking/test_watermarking.py b/tests/watermarking/test_watermarking.py index 03c01b2485..4de6022e34 100644 --- a/tests/watermarking/test_watermarking.py +++ b/tests/watermarking/test_watermarking.py @@ -26,6 +26,7 @@ from vllm.v1.watermarking.spec_decode import ( speculative_target_watermark_key, ) from vllm.v1.watermarking.watermarker import Watermarker, WatermarkSample +from vllm.v1.worker.gpu.buffer_utils import UvaBackedTensor from vllm.v1.worker.gpu.sample.sampler import Sampler from vllm.v1.worker.gpu.sample.watermark import ( draft_watermarking_mask, @@ -1499,6 +1500,53 @@ def test_draft_watermarker_deduplicates_across_a_block_on_cuda( assert not draft_logits[:, :, -1].any() +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator" +) +def test_draft_watermarker_cuda_graph_reads_current_prompt_lens(): + draft_watermarker = DraftWatermarker( + GumbelWatermarker(key=42, context_width=2), + max_num_reqs=1, + device=torch.device("cuda"), + num_speculative_steps=1, + deduplicate_contexts="single_turn", + deduplicate_contexts_max_history=None, + ) + prompt_lens = UvaBackedTensor(1, dtype=torch.int32) + contexts = torch.tensor([[5, 6]]).cuda() + enabled = torch.ones(1, dtype=torch.bool).cuda() + # [5, 6] repeats only while the prompt counts as history. + all_token_ids = torch.tensor([[5, 6, 9, 5, 6]]).cuda() + total_lens = torch.tensor([5]).cuda() + logits = torch.zeros(1, 8).cuda() + idx_mapping = torch.zeros(1, dtype=torch.int32).cuda() + temperature = torch.ones(1).cuda() + draft_step = torch.tensor(0).cuda() + + def prepare(): + draft_watermarker.prepare( + contexts, enabled, all_token_ids, prompt_lens.gpu, total_lens + ) + + def watermark_mask(): + _, mask = draft_watermarker._sampling_state( + logits, idx_mapping, temperature, draft_step, contexts + ) + return mask + + prepare() + assert not watermark_mask().item() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + mask = watermark_mask() + # A step that adds requests moves prompt_lens.gpu to the next pool buffer. + prompt_lens.np[:] = 3 + prompt_lens.copy_to_uva() + prepare() + graph.replay() + assert mask.item() + + def test_dspark_reduced_vocab_draft_sampler_applies_watermarking(monkeypatch): speculator = object.__new__(DSparkSpeculator) speculator.draft_logits = torch.zeros(2, 1, 8) diff --git a/vllm/v1/watermarking/spec_decode.py b/vllm/v1/watermarking/spec_decode.py index 36d00313ae..3cfe334a90 100644 --- a/vllm/v1/watermarking/spec_decode.py +++ b/vllm/v1/watermarking/spec_decode.py @@ -44,8 +44,8 @@ class DraftWatermarker: dtype=torch.int64, device=device, ) + self.prompt_lens = torch.zeros(max_num_reqs, dtype=torch.int32, device=device) self.all_token_ids: torch.Tensor | None = None - self.prompt_lens: torch.Tensor | None = None self.total_lens: torch.Tensor | None = None def prepare( @@ -60,7 +60,7 @@ class DraftWatermarker: self.contexts[:num_reqs].copy_(contexts) self.enabled[:num_reqs].copy_(enabled) self.all_token_ids = all_token_ids - self.prompt_lens = prompt_lens + self.prompt_lens.copy_(prompt_lens) self.total_lens = total_lens def _sampling_state( @@ -79,7 +79,6 @@ class DraftWatermarker: if self.deduplicate_contexts == "none": return steps, enabled assert self.all_token_ids is not None - assert self.prompt_lens is not None assert self.total_lens is not None enabled = draft_watermarking_mask( self.all_token_ids,