[Bugfix][Watermarking] Keep draft prompt lengths valid under CUDA graphs (#59779)

Signed-off-by: Simon Veitner <sveitner@redhat.com>
Co-authored-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com>

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)

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,