From d92a4f25f270a4c9e3a4ceb5c0fad0035e200936 Mon Sep 17 00:00:00 2001 From: Shern Shiou Tan Date: Fri, 16 Oct 2026 00:06:06 +0200 Subject: [PATCH] [watermarking] Add DFlash support to dual-key gumbel-max watermarking (#59554) Signed-off-by: Shern Shiou Tan Signed-off-by: Simon Veitner Signed-off-by: Raphael Rialland Co-authored-by: Claude Opus 5.5 Co-authored-by: Simon Veitner Co-authored-by: Raphael Rialland Co-authored-by: Lucas Wilkinson --- docs/features/watermarking.md | 6 +- tests/test_config.py | 30 ++- tests/watermarking/test_watermarking.py | 164 ++++++++++- vllm/config/vllm.py | 19 +- vllm/v1/watermarking/gumbel.py | 37 +++ vllm/v1/watermarking/spec_decode.py | 145 +++++++--- vllm/v1/watermarking/watermarker.py | 27 ++ vllm/v1/worker/gpu/sample/watermark.py | 254 +++++++++++++++++- .../gpu/spec_decode/dflash/speculator.py | 1 + vllm/v1/worker/gpu/spec_decode/speculator.py | 4 +- 10 files changed, 626 insertions(+), 61 deletions(-) diff --git a/docs/features/watermarking.md b/docs/features/watermarking.md index f51ec72eb1..adac09311c 100644 --- a/docs/features/watermarking.md +++ b/docs/features/watermarking.md @@ -125,8 +125,10 @@ generation. ## Speculative decoding Watermarking requires speculative decoding to use probabilistic draft sampling, -standard rejection sampling, and an autoregressive model-based method (`dspark`, -`eagle`, `eagle3`, or `mtp`). Parallel drafting is supported only by `dspark`. +standard rejection sampling, and a model-based method (`dflash`, `dspark`, +`eagle`, `eagle3`, or `mtp`). Parallel drafting is supported only by `dflash` +and `dspark`. DFlash drafters with a candidate head (`DFlash2DraftModel`, +`LiLiCorrDraftModel`) are not supported. A watermarking algorithm without native speculative-decoding support is rejected before model loading. Set diff --git a/tests/test_config.py b/tests/test_config.py index 6558400552..fc482bcf4e 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -3732,34 +3732,48 @@ def test_dual_key_gumbel_requires_probabilistic_drafting(): config._check_supports_watermarking() -@pytest.mark.parametrize("method", ["eagle", "eagle3", "mtp"]) -def test_dual_key_gumbel_supports_probabilistic_speculative_decoding(method): +@pytest.mark.parametrize( + ("method", "parallel_drafting"), + [ + ("eagle", False), + ("eagle3", False), + ("mtp", False), + ("dspark", True), + ("dflash", True), + ], +) +def test_dual_key_gumbel_supports_probabilistic_speculative_decoding( + method, parallel_drafting +): config = _watermarked_vllm_config() config.watermark_config = WatermarkConfig(algorithm="dual_key_gumbel", key=42) config.speculative_config = SimpleNamespace( method=method, draft_sample_method="probabilistic", rejection_sample_method="standard", - parallel_drafting=False, + parallel_drafting=parallel_drafting, ) config._check_supports_watermarking() -def test_dual_key_gumbel_supports_dspark(): +@pytest.mark.parametrize("architecture", ["DFlash2DraftModel", "LiLiCorrDraftModel"]) +def test_dual_key_gumbel_rejects_dflash_candidate_drafts(architecture): config = _watermarked_vllm_config() config.watermark_config = WatermarkConfig(algorithm="dual_key_gumbel", key=42) config.speculative_config = SimpleNamespace( - method="dspark", + method="dflash", draft_sample_method="probabilistic", rejection_sample_method="standard", parallel_drafting=True, + draft_model_config=SimpleNamespace(architectures=[architecture]), ) - config._check_supports_watermarking() + with pytest.raises(ValueError, match="candidate-head drafters"): + config._check_supports_watermarking() -def test_dual_key_gumbel_rejects_non_autoregressive_speculation(): +def test_dual_key_gumbel_rejects_unsupported_speculative_method(): config = _watermarked_vllm_config() config.watermark_config = WatermarkConfig(algorithm="dual_key_gumbel", key=42) config.speculative_config = SimpleNamespace( @@ -3769,7 +3783,7 @@ def test_dual_key_gumbel_rejects_non_autoregressive_speculation(): parallel_drafting=False, ) - with pytest.raises(ValueError, match="autoregressive model-based"): + with pytest.raises(ValueError, match="Watermarking supports only"): config._check_supports_watermarking() diff --git a/tests/watermarking/test_watermarking.py b/tests/watermarking/test_watermarking.py index 384cc30d57..eec43845db 100644 --- a/tests/watermarking/test_watermarking.py +++ b/tests/watermarking/test_watermarking.py @@ -36,6 +36,7 @@ from vllm.v1.worker.gpu.sample.watermark import ( philox_gumbel_sample, repeated_context_mask, ) +from vllm.v1.worker.gpu.spec_decode.dflash.speculator import DFlashSpeculator from vllm.v1.worker.gpu.spec_decode.dspark.speculator import DSparkSpeculator from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator @@ -254,8 +255,8 @@ def test_speculative_decoding_uses_fixed_dual_key_roles(): assert draft is not None assert draft.watermarker.prf.key == derive_watermark_key(42, b"key_a") assert draft.num_speculative_steps == 3 - assert draft.deduplicate_contexts == "all" - assert draft.deduplicate_contexts_max_history is None + assert draft.state.deduplicate_contexts == "all" + assert draft.state.deduplicate_contexts_max_history is None def test_recovery_key_rejects_the_unsplit_dual_key_watermarker(): @@ -1324,8 +1325,8 @@ def test_draft_sampler_uses_draft_key_and_advances_context(monkeypatch): deduplicate_contexts="none", deduplicate_contexts_max_history=None, ) - draft_watermarker.contexts.copy_(torch.tensor([[1, 2], [3, 4]])) - draft_watermarker.enabled = torch.tensor([True, False]) + draft_watermarker.state.contexts.copy_(torch.tensor([[1, 2], [3, 4]])) + draft_watermarker.state.enabled = torch.tensor([True, False]) speculator.draft_watermarker = draft_watermarker monkeypatch.setattr( "vllm.v1.watermarking.watermarker.gumbel_sample", @@ -1343,7 +1344,7 @@ def test_draft_sampler_uses_draft_key_and_advances_context(monkeypatch): ) assert torch.equal(sampled, torch.tensor([7, 4])) - assert torch.equal(draft_watermarker.contexts, torch.tensor([[2, 7], [4, 4]])) + assert torch.equal(draft_watermarker.state.contexts, torch.tensor([[2, 7], [4, 4]])) def test_draft_sampler_uses_ordinary_samples_for_repeated_contexts(monkeypatch): @@ -1602,11 +1603,12 @@ def test_draft_watermarker_deduplicates_across_a_block_on_cuda( [True, True], ] assert torch.equal( - draft_watermarker.prior_contexts[:num_reqs].cpu(), + draft_watermarker.state.prior_contexts[:num_reqs].cpu(), torch.tensor([[[5, 6], [6, 6], [6, 6]], [[6, 3], [3, 6], [6, 6]]]), ) assert torch.equal( - draft_watermarker.contexts[:num_reqs].cpu(), torch.tensor([[6, 6], [6, 6]]) + draft_watermarker.state.contexts[:num_reqs].cpu(), + torch.tensor([[6, 6], [6, 6]]), ) for step, logits in enumerate(step_logits): # The 0-d draft step selects the logits cache column. @@ -1661,6 +1663,154 @@ def test_draft_watermarker_cuda_graph_reads_current_prompt_lens(): assert mask.item() +def test_dflash_draft_sampler_watermarks_drafts_in_step_order(monkeypatch): + class NextTokenWatermarker: + context_width = 2 + + @staticmethod + def _try_sample_block(*args, **kwargs): + return None + + @staticmethod + def sample(logits, contexts, random_sampler=None, skip_mask=None): + watermarked = contexts[:, -1] + 1 + return WatermarkSample( + torch.where(skip_mask, random_sampler(logits), watermarked), logits + ) + + def ordinary_sample(logits, idx_mapping, temperature, seed, pos, **kwargs): + assert kwargs["use_fp64"] + cache, col = kwargs["logits_cache"], kwargs["logits_cache_col"] + cache[idx_mapping, col] = kwargs["logits_cache_source"] + return seed[idx_mapping] + pos + + monkeypatch.setattr( + "vllm.v1.watermarking.watermarker.gumbel_sample", ordinary_sample + ) + num_reqs, num_steps, vocab_size = 3, 2, 4 + draft_watermarker = DraftWatermarker( + NextTokenWatermarker(), + max_num_reqs=8, + device=torch.device("cpu"), + num_speculative_steps=num_steps, + deduplicate_contexts="none", + deduplicate_contexts_max_history=None, + ) + # Batch rows hold request slots 2, 0, 1. Slot 0 is greedy and row 2 opted out. + draft_watermarker.state.contexts[:num_reqs] = torch.tensor([[1, 2], [3, 4], [5, 6]]) + draft_watermarker.state.enabled[:num_reqs] = torch.tensor([True, True, False]) + speculator = object.__new__(DFlashSpeculator) + speculator.model = SimpleNamespace( + compute_logits=lambda hidden_states: hidden_states + ) + speculator.draft_watermarker = draft_watermarker + speculator.use_fp64_gumbel = True + speculator.acceptance_estimator = None + batch_slots = torch.tensor([2, 0, 1]) + logits = torch.arange(num_reqs * num_steps * vocab_size, dtype=torch.float32) + draft_logits = torch.zeros(num_reqs, num_steps, vocab_size) + + sampled = speculator.sample_draft( + hidden_states=logits.view(-1, vocab_size), + sample_src_positions=torch.arange(num_reqs * num_steps), + idx_mapping=batch_slots.repeat_interleave(num_steps), + temperature=torch.tensor([0.0, 1.0, 0.5]), + seeds=torch.tensor([100, 200, 300]), + draft_step=torch.arange(num_steps).repeat(num_reqs), + draft_logits=draft_logits, + num_steps=num_steps, + ) + + assert torch.equal(sampled, torch.tensor([3, 4, 102, 103, 204, 205])) + assert torch.equal(draft_logits[batch_slots], logits.view(num_reqs, num_steps, -1)) + + +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="requires a CUDA-like accelerator" +) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("deduplicate_contexts", ["none", "single_turn", "all"]) +def test_dflash_draft_sampler_matches_step_loop_on_cuda(deduplicate_contexts, dtype): + """The fused CUDA `sample_block` is bit-exact with sampling step by step. + + Covers repeated contexts, a NaN logit, and greedy, opted-out and padding + rows; padding rows are discarded, so only valid rows are compared. + """ + num_reqs, num_steps, vocab_size, max_num_reqs = 4, 5, 1500, 6 + generator = torch.Generator().manual_seed(0) + logits = torch.randn(num_reqs, num_steps, vocab_size, generator=generator) * 3 + logits[..., :3] += 6 + logits[0, 1, 5] = float("nan") + logits = logits.to(dtype).cuda() + idx_mapping = torch.tensor([4, 2, 0, -1], dtype=torch.int32) + valid = (idx_mapping >= 0).cuda() + idx_mapping = idx_mapping.repeat_interleave(num_steps).cuda() + temperature = torch.tensor([0.7, 1.0, 0.0, 1.3, 0.6, 0.9]).cuda() + seeds = torch.arange(max_num_reqs).cuda() * 1000 + pos = (torch.arange(num_steps).repeat(num_reqs) + 20).cuda() + draft_step = torch.arange(num_steps).repeat(num_reqs).cuda() + contexts = torch.randint(0, 3, (num_reqs, 2), generator=generator) + contexts[0, 0] = -1 + all_token_ids = torch.randint(0, 3, (max_num_reqs, 32), generator=generator) + + def run(fused: bool): + draft_watermarker = DraftWatermarker( + GumbelWatermarker(key=42, context_width=2), + max_num_reqs=max_num_reqs, + device=torch.device("cuda"), + num_speculative_steps=num_steps, + deduplicate_contexts=deduplicate_contexts, + deduplicate_contexts_max_history=None, + ) + draft_watermarker.prepare( + contexts=contexts.cuda(), + enabled=torch.tensor([True, False, True, True]).cuda(), + all_token_ids=all_token_ids.cuda(), + prompt_lens=torch.full((max_num_reqs,), 4).cuda(), + total_lens=torch.full((max_num_reqs,), 24).cuda(), + ) + draft_logits = torch.zeros(max_num_reqs, num_steps, vocab_size, dtype=dtype) + draft_logits = draft_logits.cuda() + args = (temperature, seeds) + if fused: + sampled = draft_watermarker.sample_block( + logits.view(-1, vocab_size), + idx_mapping, + *args, + pos, + logits_cache=draft_logits, + logits_cache_col=draft_step, + ) + else: + by_step = [t.view(num_reqs, num_steps) for t in (idx_mapping, pos)] + steps = draft_step.view(num_reqs, num_steps) + sampled = torch.stack( + [ + draft_watermarker.sample( + logits[:, step], + by_step[0][:, step], + *args, + by_step[1][:, step], + logits_cache=draft_logits, + logits_cache_col=steps[:, step], + apply_temperature=True, + is_drafting=True, + ) + for step in range(num_steps) + ], + dim=1, + ).view(-1) + return ( + sampled.view(num_reqs, num_steps)[valid], + draft_watermarker.state.contexts[:num_reqs][valid], + draft_watermarker.state.prior_contexts, + draft_logits, + ) + + for fused, reference in zip(run(fused=True), run(fused=False)): + torch.testing.assert_close(fused, reference, rtol=0, atol=0, equal_nan=True) + + 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/config/vllm.py b/vllm/config/vllm.py index 608652e08c..3345d89fbd 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -1318,15 +1318,26 @@ class VllmConfig: ) if ( speculative_config.parallel_drafting - and speculative_config.method != "dspark" + and speculative_config.method not in ("dflash", "dspark") ): raise ValueError( "Parallel speculative drafting is not supported with watermarking." ) - if speculative_config.method not in ("dspark", "eagle", "eagle3", "mtp"): + if speculative_config.method not in ( + "dflash", + "dspark", + "eagle", + "eagle3", + "mtp", + ): raise ValueError( - "Watermarking supports only autoregressive model-based " - "speculative decoding." + "Watermarking supports only dflash, dspark, eagle, eagle3 and " + "mtp speculative decoding." + ) + if self._is_dflash_candidate_draft(): + raise ValueError( + "Watermarking does not support DFlash2 and LiLiCorr " + "candidate-head drafters yet." ) if ( not watermark_config.allow_target_only_watermarking diff --git a/vllm/v1/watermarking/gumbel.py b/vllm/v1/watermarking/gumbel.py index 2f953dd389..31e7bdd616 100644 --- a/vllm/v1/watermarking/gumbel.py +++ b/vllm/v1/watermarking/gumbel.py @@ -14,6 +14,7 @@ from vllm.v1.watermarking.detector import ( ) from vllm.v1.watermarking.prfs import PhiloxPRF, WatermarkPRF, create_prf from vllm.v1.watermarking.watermarker import ( + DraftContextState, RandomSampler, SupportsSpeculativeDecoding, Watermarker, @@ -101,6 +102,42 @@ class GumbelWatermarker(Watermarker): ) return WatermarkSample(token_ids, logits) + def _try_sample_block( + self, + logits: torch.Tensor, + num_steps: int, + random_sampler: RandomSampler, + state: DraftContextState, + ) -> torch.Tensor | None: + if type(self.prf) is not PhiloxPRF or logits.device.type != "cuda": + return None + + from vllm.v1.worker.gpu.sample.watermark import draft_philox_gumbel_sample + + assert random_sampler.logits_cache is not None + assert random_sampler.logits_cache_col is not None + return draft_philox_gumbel_sample( + logits, + state.contexts, + self.prf.key, + num_steps=num_steps, + expanded_idx_mapping=random_sampler.expanded_idx_mapping, + temperatures=random_sampler.temperatures, + seeds=random_sampler.seeds, + positions=random_sampler.positions, + enabled=state.enabled, + logits_cache=random_sampler.logits_cache, + logits_cache_col=random_sampler.logits_cache_col, + use_fp64=random_sampler.use_fp64, + prior_contexts=state.prior_contexts, + all_token_ids=state.all_token_ids, + prompt_lens=state.prompt_lens, + total_lens=state.total_lens, + deduplicate=state.deduplicate_contexts != "none", + max_history=state.deduplicate_contexts_max_history, + include_prompt=state.deduplicate_contexts == "all", + ) + class DualKeyGumbelWatermarker(Watermarker, SupportsSpeculativeDecoding): def __init__( diff --git a/vllm/v1/watermarking/spec_decode.py b/vllm/v1/watermarking/spec_decode.py index 3cfe334a90..6c432db959 100644 --- a/vllm/v1/watermarking/spec_decode.py +++ b/vllm/v1/watermarking/spec_decode.py @@ -8,6 +8,7 @@ from vllm.v1.watermarking.factory import create_watermarker from vllm.v1.watermarking.gumbel import GumbelWatermarker from vllm.v1.watermarking.prfs import PhiloxPRF from vllm.v1.watermarking.watermarker import ( + DraftContextState, RandomSampler, SupportsSpeculativeDecoding, Watermarker, @@ -28,25 +29,25 @@ class DraftWatermarker: ) -> None: self.watermarker = watermarker self.num_speculative_steps = num_speculative_steps - self.deduplicate_contexts = deduplicate_contexts - self.deduplicate_contexts_max_history = deduplicate_contexts_max_history - self.contexts = torch.zeros( - max_num_reqs, - watermarker.context_width, - dtype=torch.int64, - device=device, + self.state = DraftContextState( + contexts=torch.zeros( + max_num_reqs, + watermarker.context_width, + dtype=torch.int64, + device=device, + ), + enabled=torch.zeros(max_num_reqs, dtype=torch.bool, device=device), + prior_contexts=torch.zeros( + max_num_reqs, + num_speculative_steps, + watermarker.context_width, + dtype=torch.int64, + device=device, + ), + prompt_lens=torch.zeros(max_num_reqs, dtype=torch.int32, device=device), + deduplicate_contexts=deduplicate_contexts, + deduplicate_contexts_max_history=deduplicate_contexts_max_history, ) - self.enabled = torch.zeros(max_num_reqs, dtype=torch.bool, device=device) - self.prior_contexts = torch.zeros( - max_num_reqs, - num_speculative_steps, - watermarker.context_width, - 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.total_lens: torch.Tensor | None = None def prepare( self, @@ -57,11 +58,12 @@ class DraftWatermarker: total_lens: torch.Tensor, ) -> None: num_reqs = contexts.shape[0] - self.contexts[:num_reqs].copy_(contexts) - self.enabled[:num_reqs].copy_(enabled) - self.all_token_ids = all_token_ids - self.prompt_lens.copy_(prompt_lens) - self.total_lens = total_lens + state = self.state + state.contexts[:num_reqs].copy_(contexts) + state.enabled[:num_reqs].copy_(enabled) + state.all_token_ids = all_token_ids + state.prompt_lens.copy_(prompt_lens) + state.total_lens = total_lens def _sampling_state( self, @@ -75,22 +77,23 @@ class DraftWatermarker: steps = torch.as_tensor(draft_step, device=logits.device, dtype=torch.int64) if steps.ndim == 0: steps = steps.expand(num_rows) - enabled = self.enabled[:num_rows] & (request_temperatures != 0) - if self.deduplicate_contexts == "none": + state = self.state + enabled = state.enabled[:num_rows] & (request_temperatures != 0) + if state.deduplicate_contexts == "none": return steps, enabled - assert self.all_token_ids is not None - assert self.total_lens is not None + assert state.all_token_ids is not None + assert state.total_lens is not None enabled = draft_watermarking_mask( - self.all_token_ids, + state.all_token_ids, idx_mapping, - self.prompt_lens, - self.total_lens, - self.prior_contexts, + state.prompt_lens, + state.total_lens, + state.prior_contexts, contexts, steps, enabled, - self.deduplicate_contexts_max_history, - include_prompt=self.deduplicate_contexts == "all", + state.deduplicate_contexts_max_history, + include_prompt=state.deduplicate_contexts == "all", ) return steps, enabled @@ -107,14 +110,26 @@ class DraftWatermarker: logits_cache_col: torch.Tensor | None = None, use_fp64: bool = False, logits_cache_source: torch.Tensor | None = None, + num_steps: int = 1, ) -> torch.Tensor: + if num_steps > 1: + return self.sample_block( + logits, + idx_mapping, + temperature, + seed, + pos, + logits_cache=logits_cache, + logits_cache_col=logits_cache_col, + use_fp64=use_fp64, + ) assert apply_temperature assert is_drafting assert logits_cache is not None assert logits_cache_col is not None num_rows = logits.shape[0] request_temperatures = temperature[idx_mapping] - contexts = self.contexts[:num_rows] + contexts = self.state.contexts[:num_rows] steps, enabled = self._sampling_state( logits, idx_mapping, request_temperatures, logits_cache_col, contexts ) @@ -145,6 +160,68 @@ class DraftWatermarker: contexts.copy_(torch.cat((contexts[:, 1:], sampled.unsqueeze(-1)), dim=-1)) return sampled + def sample_block( + self, + logits: torch.Tensor, + idx_mapping: torch.Tensor, + temperature: torch.Tensor, + seed: torch.Tensor, + pos: torch.Tensor, + logits_cache: torch.Tensor | None = None, + logits_cache_col: torch.Tensor | None = None, + use_fp64: bool = False, + ) -> torch.Tensor: + """Sample drafts flattened as (request, step), one step at a time. + The target verifies each draft with a context that includes the earlier + drafts of its request, so steps must be sampled in order even when + their logits were computed in parallel. + """ + assert logits_cache_col is not None + num_steps = self.num_speculative_steps + sampled_block = self.watermarker._try_sample_block( + logits, + num_steps, + RandomSampler( + expanded_idx_mapping=idx_mapping, + temperatures=temperature, + seeds=seed, + positions=pos, + use_fp64=use_fp64, + is_drafting=True, + logits_cache=logits_cache, + logits_cache_col=logits_cache_col, + ), + self.state, + ) + if sampled_block is not None: + return sampled_block + num_reqs = logits.shape[0] // num_steps + + def by_step(tensor: torch.Tensor) -> torch.Tensor: + return tensor.view(num_reqs, num_steps, *tensor.shape[1:]) + + logits = by_step(logits) + idx_mapping = by_step(idx_mapping) + pos = by_step(pos) + logits_cache_col = by_step(logits_cache_col) + sampled = torch.empty( + num_reqs, num_steps, dtype=torch.int64, device=logits.device + ) + for step in range(num_steps): + sampled[:, step] = self.sample( + logits[:, step], + idx_mapping[:, step], + temperature, + seed, + pos[:, step], + apply_temperature=True, + is_drafting=True, + logits_cache=logits_cache, + logits_cache_col=logits_cache_col[:, step], + use_fp64=use_fp64, + ) + return sampled.view(-1) + def create_speculative_target_watermarker(watermarker: Watermarker) -> Watermarker: if isinstance(watermarker, SupportsSpeculativeDecoding): diff --git a/vllm/v1/watermarking/watermarker.py b/vllm/v1/watermarking/watermarker.py index 15917d2520..13b806faac 100644 --- a/vllm/v1/watermarking/watermarker.py +++ b/vllm/v1/watermarking/watermarker.py @@ -7,6 +7,7 @@ from typing import Protocol, runtime_checkable import torch +from vllm.config.watermarking import WatermarkContextScope from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample @@ -44,6 +45,18 @@ class RandomSampler: ) +@dataclass +class DraftContextState: + contexts: torch.Tensor + enabled: torch.Tensor + prior_contexts: torch.Tensor + prompt_lens: torch.Tensor + deduplicate_contexts: WatermarkContextScope + deduplicate_contexts_max_history: int | None + all_token_ids: torch.Tensor | None = None + total_lens: torch.Tensor | None = None + + class Watermarker(ABC): @property @abstractmethod @@ -94,6 +107,20 @@ class Watermarker(ABC): ) -> WatermarkSample | None: return None + def _try_sample_block( + self, + logits: torch.Tensor, + num_steps: int, + random_sampler: RandomSampler, + state: DraftContextState, + ) -> torch.Tensor | None: + """Sample a (request, step) draft block in one fused pass. + + Returns None when there is no fused implementation, in which case the + caller samples the block one step at a time. + """ + return None + @runtime_checkable class SupportsSpeculativeDecoding(Protocol): diff --git a/vllm/v1/worker/gpu/sample/watermark.py b/vllm/v1/worker/gpu/sample/watermark.py index d366b7622a..00a137fd1b 100644 --- a/vllm/v1/worker/gpu/sample/watermark.py +++ b/vllm/v1/worker/gpu/sample/watermark.py @@ -530,8 +530,10 @@ def _philox_gumbel_from_logits(logits, word): @triton.jit -def _gumbel_value(logits_ptr, output, mask): +def _gumbel_value(logits_ptr, output, mask, temperature): logits = tl.load(logits_ptr, mask=mask, other=float("-inf")).to(tl.float32) + if temperature is not None: + logits = tl_math.div_rn(logits, temperature) return _philox_gumbel_from_logits(logits, output) @@ -606,6 +608,7 @@ def _philox_gumbel_kernel( USE_FP64: tl.constexpr, IS_DRAFTING: tl.constexpr, PER_TOKEN_COL: tl.constexpr, + APPLY_TEMPERATURE: tl.constexpr, ): row = tl.program_id(0).to(tl.int64) block_index = tl.program_id(1) @@ -619,6 +622,13 @@ def _philox_gumbel_kernel( req_state_idx = tl.load(expanded_idx_mapping_ptr + row).to(tl.int64) valid_req = req_state_idx >= 0 + row_temperature = None + if APPLY_TEMPERATURE: + row_temperature = tl.load( + temp_ptr + req_state_idx, mask=valid_req, other=0.0 + ).to(tl.float32) + row_temperature = tl.where(row_temperature == 0, 1.0, row_temperature) + if logits_cache_ptr is not None: if PER_TOKEN_COL: col = tl.load(logits_cache_col_ptr + row) @@ -650,6 +660,8 @@ def _philox_gumbel_kernel( logits = tl.load(logits_row + candidate, mask=mask, other=float("-inf")).to( tl.float32 ) + if row_temperature is not None: + logits = tl_math.div_rn(logits, row_temperature) value, index = gumbel_noised_argmax( logits, candidate, @@ -683,16 +695,16 @@ def _philox_gumbel_kernel( candidate_3 = candidate_0 + 3 logits_row = logits_ptr + row * logits_stride value_0 = _gumbel_value( - logits_row + candidate_0, output_0, candidate_0 < vocab_size + logits_row + candidate_0, output_0, candidate_0 < vocab_size, row_temperature ) value_1 = _gumbel_value( - logits_row + candidate_1, output_1, candidate_1 < vocab_size + logits_row + candidate_1, output_1, candidate_1 < vocab_size, row_temperature ) value_2 = _gumbel_value( - logits_row + candidate_2, output_2, candidate_2 < vocab_size + logits_row + candidate_2, output_2, candidate_2 < vocab_size, row_temperature ) value_3 = _gumbel_value( - logits_row + candidate_3, output_3, candidate_3 < vocab_size + logits_row + candidate_3, output_3, candidate_3 < vocab_size, row_temperature ) best_value = value_0 best_token = candidate_0 @@ -799,6 +811,238 @@ def philox_gumbel_sample( USE_FP64=use_fp64, IS_DRAFTING=is_drafting, PER_TOKEN_COL=logits_cache_col is not None and logits_cache_col.dim() > 0, + APPLY_TEMPERATURE=False, ) max_block_index = local_max.argmax(dim=-1, keepdim=True) return local_argmax.gather(dim=-1, index=max_block_index).view(-1) + + +@triton.jit(do_not_specialize=["step"]) +def _draft_step_kernel( + sampled_ptr, + sampled_stride, + step, + local_argmax_ptr, + local_max_ptr, + num_blocks, + contexts_ptr, + context_stride, + skip_mask_ptr, + next_req_indices_ptr, + next_steps_ptr, + temp_ptr, + enabled_ptr, + all_token_ids_ptr, + all_token_ids_stride, + prompt_lens_ptr, + total_lens_ptr, + prior_contexts_ptr, + prior_contexts_stride_0, + prior_contexts_stride_1, + CONTEXT_WIDTH: tl.constexpr, + NUM_SPECULATIVE_STEPS: tl.constexpr, + MAX_HISTORY: tl.constexpr, + INCLUDE_PROMPT: tl.constexpr, + DEDUPLICATE: tl.constexpr, + SAMPLE: tl.constexpr, + PREPARE_NEXT: tl.constexpr, + BLOCKS_POW2: tl.constexpr, +): + """Finish draft step `step` (if SAMPLE) and build the next step's skip mask.""" + row = tl.program_id(0).to(tl.int64) + contexts_row_ptr = contexts_ptr + row * context_stride + if SAMPLE: + # Same pick as local_max.argmax(dim=-1): the first NaN, else the first max. + blocks = tl.arange(0, BLOCKS_POW2) + values = tl.load( + local_max_ptr + row * num_blocks + blocks, + mask=blocks < num_blocks, + other=float("-inf"), + ) + is_nan = values != values + best = tl.max(tl.where(is_nan, float("-inf"), values), axis=0) + first_nan = tl.min(tl.where(is_nan, blocks, BLOCKS_POW2), axis=0) + first_max = tl.min( + tl.where((values == best) & (blocks < num_blocks), blocks, BLOCKS_POW2), + axis=0, + ) + block = tl.where(first_nan < BLOCKS_POW2, first_nan, first_max) + token = tl.load(local_argmax_ptr + row * num_blocks + block) + tl.store(sampled_ptr + row * sampled_stride + step, token) + + offsets = tl.arange(0, BLOCKS_POW2) + context = tl.load( + contexts_row_ptr + offsets + 1, + mask=offsets < CONTEXT_WIDTH - 1, + other=token, + ) + tl.debug_barrier() + tl.store(contexts_row_ptr + offsets, context, mask=offsets < CONTEXT_WIDTH) + tl.debug_barrier() + + if PREPARE_NEXT: + req_idx = tl.load(next_req_indices_ptr + row).to(tl.int64) + temperature = tl.load(temp_ptr + req_idx, mask=req_idx >= 0, other=0.0) + active = tl.load(enabled_ptr + row) & (temperature != 0) + if DEDUPLICATE: + # The dedup kernel body runs in place on the skip mask buffer. + tl.store(skip_mask_ptr + row, active) + tl.debug_barrier() + _repeated_context_mask_kernel( + skip_mask_ptr, + all_token_ids_ptr, + all_token_ids_stride, + next_req_indices_ptr, + prompt_lens_ptr, + total_lens_ptr, + next_steps_ptr, + 1, + contexts_ptr, + context_stride, + None, + prior_contexts_ptr, + prior_contexts_stride_0, + prior_contexts_stride_1, + next_steps_ptr, + 1, + skip_mask_ptr, + CONTEXT_WIDTH=CONTEXT_WIDTH, + NUM_SPECULATIVE_STEPS=NUM_SPECULATIVE_STEPS, + MAX_HISTORY=MAX_HISTORY, + INCLUDE_PROMPT=INCLUDE_PROMPT, + SKIP_PARTIAL_CONTEXT=False, + HAS_HISTORY_OFFSETS=True, + SPECULATIVE_CONTEXTS=False, + DRAFT_CONTEXTS=True, + BLOCK=512, + ) + tl.debug_barrier() + active = tl.load(skip_mask_ptr + row) + tl.store(skip_mask_ptr + row, active == 0) + + +def draft_philox_gumbel_sample( + logits: torch.Tensor, + contexts: torch.Tensor, + key: int, + *, + num_steps: int, + expanded_idx_mapping: torch.Tensor, + temperatures: torch.Tensor, + seeds: torch.Tensor, + positions: torch.Tensor, + enabled: torch.Tensor, + logits_cache: torch.Tensor, + logits_cache_col: torch.Tensor, + use_fp64: bool, + prior_contexts: torch.Tensor, + all_token_ids: torch.Tensor | None, + prompt_lens: torch.Tensor, + total_lens: torch.Tensor | None, + deduplicate: bool, + max_history: int | None, + include_prompt: bool, +) -> torch.Tensor: + """Sample `num_steps` draft steps flattened as (request, step), in step order. + + Bit-exact with sampling each step by `DraftWatermarker.sample`, but with two + kernel launches per step: the Philox block argmax (with the temperature + divide and the logits cache write), then a per-row kernel that picks the + token, shifts `contexts` and builds the next step's skip mask. Padding rows + (index -1) read temperature 0 and are never watermarked, so they may differ. + """ + num_rows, vocab_size = logits.shape + num_reqs = num_rows // num_steps + assert logits.stride(-1) == 1 and contexts.stride(-1) == 1 + assert logits.dtype == logits_cache.dtype, ( + "logits cache source and destination must have the same dtype" + ) + assert logits_cache.size(-1) >= vocab_size, ( + f"draft logits cache vocab dim ({logits_cache.size(-1)}) is narrower " + f"than the sampled logits ({vocab_size}). Cached logits would be truncated." + ) + assert not deduplicate or (all_token_ids is not None and total_lens is not None), ( + "context deduplication requires all_token_ids and total_lens" + ) + logits = logits.view(num_reqs, num_steps, vocab_size) + # Step-major copies, so each step reads contiguous per-request slices. + expanded_idx_mapping, positions, logits_cache_col = ( + tensor.view(num_reqs, num_steps).t().contiguous() + for tensor in (expanded_idx_mapping, positions, logits_cache_col) + ) + block_size = 1024 + num_blocks = triton.cdiv(vocab_size, block_size) + local_argmax = logits.new_empty(num_reqs, num_blocks, dtype=torch.int64) + local_max = logits.new_empty( + num_reqs, num_blocks, dtype=torch.float64 if use_fp64 else torch.float32 + ) + skip_mask = torch.empty(num_reqs, dtype=torch.bool, device=logits.device) + sampled = logits.new_empty(num_reqs, num_steps, dtype=torch.int64) + + def finish_step(step: int) -> None: + next_step = min(step + 1, num_steps - 1) + _draft_step_kernel[(num_reqs,)]( + sampled, + sampled.stride(0), + step, + local_argmax, + local_max, + num_blocks, + contexts, + contexts.stride(0), + skip_mask, + expanded_idx_mapping[next_step], + logits_cache_col[next_step], + temperatures, + enabled, + all_token_ids, + 0 if all_token_ids is None else all_token_ids.stride(0), + prompt_lens, + total_lens, + prior_contexts, + prior_contexts.stride(0), + prior_contexts.stride(1), + CONTEXT_WIDTH=contexts.shape[-1], + NUM_SPECULATIVE_STEPS=prior_contexts.shape[1], + MAX_HISTORY=0 if max_history is None else max_history, + INCLUDE_PROMPT=include_prompt, + DEDUPLICATE=deduplicate, + SAMPLE=step >= 0, + PREPARE_NEXT=step + 1 < num_steps, + BLOCKS_POW2=triton.next_power_of_2(max(num_blocks, contexts.shape[-1])), + ) + + finish_step(-1) + for step in range(num_steps): + _philox_gumbel_kernel[(num_reqs, num_blocks)]( + local_argmax, + local_argmax.stride(0), + local_max, + local_max.stride(0), + logits[:, step], + logits.stride(0), + contexts, + contexts.stride(0), + skip_mask, + expanded_idx_mapping[step], + seeds, + positions[step], + temperatures, + logits_cache, + logits_cache.stride(0), + logits_cache.stride(1), + logits_cache_col[step], + logits[:, step], + logits.stride(0), + key & _UINT32_MASK_VALUE, + key >> 32, + vocab_size, + CONTEXT_WIDTH=contexts.shape[-1], + BLOCK_SIZE=block_size, + USE_FP64=use_fp64, + IS_DRAFTING=True, + PER_TOKEN_COL=True, + APPLY_TEMPERATURE=True, + ) + finish_step(step) + return sampled.view(-1) diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index a75e84100e..af519bb5db 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -366,6 +366,7 @@ class DFlashSpeculator(DraftModelSpeculator): self.seeds, self.sample_col[:num_sample], self.draft_logits, + num_steps=self.num_speculative_steps, ) self.draft_tokens[:num_reqs] = draft_tokens.view( num_reqs, self.num_speculative_steps diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index d86a0b0351..c54df5ac17 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/speculator.py @@ -3,6 +3,7 @@ from abc import ABC, abstractmethod from collections.abc import Mapping from dataclasses import replace +from functools import partial from typing import TYPE_CHECKING, Any import numpy as np @@ -446,6 +447,7 @@ class DraftModelSpeculator(BaseSpeculator): draft_step: torch.Tensor, draft_logits: torch.Tensor | None, spec_step_idx: int = 0, + num_steps: int = 1, ) -> torch.Tensor: if draft_logits is None and self.use_local_argmax_reduction: return self.get_draft_top_tokens(hidden_states, spec_step_idx) @@ -455,7 +457,7 @@ class DraftModelSpeculator(BaseSpeculator): sampler = ( gumbel_sample if self.draft_watermarker is None - else self.draft_watermarker.sample + else partial(self.draft_watermarker.sample, num_steps=num_steps) ) sampled = sampler( logits,