iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[watermarking] Add DFlash support to dual-key gumbel-max watermarking (#59554)

Browse Source 
Signed-off-by: Shern Shiou Tan <shernshiou@gmail.com>
Signed-off-by: Simon Veitner <sveitner@redhat.com>
Signed-off-by: Raphael Rialland <raphael.rialland@mistral.ai>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Co-authored-by: Simon Veitner <sveitner@redhat.com>
Co-authored-by: Raphael Rialland <raphael.rialland@mistral.ai>
Co-authored-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com>
This commit is contained in:
Shern Shiou Tan
2026-10-15 22:06:06 +00:00
committed by GitHub
co-authored by Claude Opus 5.5 Simon Veitner Raphael Rialland Lucas Wilkinson
parent e711723352
commit d92a4f25f2
10 changed files with 626 additions and 61 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/features/watermarking.md
+4 -2
View File
@@ -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
tests/test_config.py
+22 -8
View File
@@ -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()
tests/watermarking/test_watermarking.py
+157 -7
View File
@@ -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)
vllm/config/vllm.py
+15 -4
View File
@@ -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
vllm/v1/watermarking/gumbel.py
+37
View File
@@ -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__(
vllm/v1/watermarking/spec_decode.py
+111 -34
View File
@@ -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):
vllm/v1/watermarking/watermarker.py
+27
View File
@@ -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):
vllm/v1/worker/gpu/sample/watermark.py
+249 -5
View File
@@ -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)
vllm/v1/worker/gpu/spec_decode/dflash/speculator.py
+1
View File
@@ -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
vllm/v1/worker/gpu/spec_decode/speculator.py
+3 -1
View File
@@ -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,