LLM Inference Series: How LLMs Embed Watermarks Without Distorting Token Distributions

In this post, we will cover in depth how LLM watermarking works, including both embedding and detection. We will focus on Gumbel-Max Watermarking (which does not impact LLM training and alters only the sampling procedure) and how it enables lightweight detection without access to the model.

We start by presenting the theory, followed by some proofs, and then look at vLLM implementation.

Why would we care about watermarking? Well, there are many reasons including identifying synthetic text, detecting bot-generated propaganda, excluding LLM-generated output from future training data [1,4], complying with regulations that require AI-generated content to be identifiable [2] and detecting unauthorized distillation [6].

LLM Training, Inference and Sampling

During LLM training, the token distribution, or more precisely the conditional token distribution \(p_\theta(x_i \mid x_{<i})\), is a categorical distribution over the tokens, parameterized using the output of the deep learning model, that is, the hidden states returned by the last transformer layer, which are then projected into logits by the output head (the parameters of the model are shared to produce these outputs for all tokens in the input stream, and each token's hidden state basically compresses information from the previous context due to how attention works). So the model learns how to parameterize this distribution, and then, once we have a categorical distribution at hand, we can sample tokens. The goal of training is to learn how to obtain this conditional distribution.

Now, during LLM inference, each request is scheduled for one or more steps to complete the prefill and populate the KV cache entries. The hidden state of the last position is then used to predict the next token (for a given request), again using the output logits, which parameterize the next-token distribution. Sampling then continues until a stopping condition is triggered.

During decoding, only one new token is scheduled per request at each step, and each step produces one token when speculative decoding is not enabled. When speculative decoding is enabled, a draft model is employed to produce draft tokens, which are then verified by the target model using a rejection-sampling procedure to accept or reject speculative tokens and potentially sample a bonus token (yay).

We will cover the internal details of LLM inference engines, down to the kernel level, very soon, but for now, let's focus on the topic of this post, watermarking.

Just remember that we have a categorical distribution parameterized by the underlying foundation model, which is used to sample the next token. This categorical distribution is different for each context, \(p_\theta(x_i \mid x_{<i})\), and the model has learned how to produce it.

Introducing LLM Watermarking

Basically, the goal of watermarking is to change the original sampling scheme (i.e., altering the LLM decoding process) by embedding a hidden watermark into the generated token stream, i.e., the output tokens. Given a sequence of tokens, one can then use a well-defined mechanism to statistically detect the presence of the watermark (more on that later).

Aaronson et al. proposed sampling a random vector \(R \in [0,1]^{|V|}\) depending on some secret key, and then choosing the next token using \(x_t = \arg\max_k R_k^{1/p_k}\), where \(p_k\) is the probability of sampling token \(k\). The random variable \(\arg\max_k R_k^{1/p_k}\) is thus defined over the vocabulary (intuitively this sampling scheme encourages the generation of tokens that have a high \(R_k\) value; remember that!).

As we will show later, it also presents an interesting property:

\(P(\arg\max_k R_k^{1/p_k} = m) = p_m\)

That is, the probability of generating a token is not altered on expectation over the secret key (we get the same LLM’s original distribution).

When the watermarking scheme is not in place, the vector \(R\) is assumed to follow a uniform distribution. On the other hand, when the watermarking scheme is in place, the vector \(R\) is generated using a pseudorandom function that depends on the context, i.e., a window of \(h\) previous tokens, and a secret key. More precisely, a secret-keyed cryptographic hash function is applied to the previous \(h\) tokens, \((x_{-h},\ldots,x_{-1})\), to produce a seed, which is then used to initialize a random number generator.

 To detect whether a piece of text is watermarked, the \(R_t\) value corresponding to each observed token \(x_t\) is reconstructed using the secret key and its context. These values are collected and used to produce a score:

\(S_T = -\sum_{t=1}^{T}\ln(1-R_t)\)

Because the score is expected to be higher when the watermarking scheme is in place (depending on the entropy and confidence of the model; larger models tend to be more confident and thus lower entropy and RLHF can also reduce entropy [10]), we can use a statistical hypothesis test to detect the watermark.

More precisely, if

\(\operatorname{p\text{-}value}(s) = P(S_T \geq s \mid H_0) < \mathrm{FPR}\)

where \(\mathrm{FPR}\) is the chosen false-positive-rate threshold, we conclude that there is statistical evidence of a watermark. In other words, if the p-value is small, the observed score is unlikely under \(H_0\), where the score follows a \(\Gamma(T,1)\) distribution, and thus the watermarking scheme is believed to have been involved. Here, \(H_0\) means that the text was generated independently of the watermarking key (non-watermarked text), whereas \(H_1\) means that the text was generated using the watermarking scheme.

As we will later, under watermarking, i.e., under \(H_1\), the score is expected to be higher.

The Gumbel-Max Trick

Let \(p=(p_1,p_2,\ldots,p_K)\) be the parameters of a categorical distribution, that is,

\[\Pr(X=i)=p_i,\qquad \sum_{i=1}^K p_i=1\]

The Gumbel-Max trick states that we can draw samples \(X\sim p\) using

\[ X\overset{d}{=}\underset{k=1,\ldots,K}{\arg\max}\;\bigl(g_k+\log p_k\bigr),\]

where the \(g_k\) are independent samples drawn from the \(\operatorname{Gumbel}(0,1)\) distribution. We assume \(p_k>0\) for the tokens considered.

To sample from \(\operatorname{Gumbel}(0,1)\), we sample \(U\sim\operatorname{Uniform}(0,1)\) and set

\[ g=-\log(-\log U)\]

The PDF and CDF of \(\operatorname{Gumbel}(0,1)\) are, respectively,

\[f_g(x)=e^{-x}e^{-e^{-x}},\qquad F_g(x)=e^{-e^{-x}}\]

Now let us prove that

\[\Pr\!\left(\underset{k}{\arg\max}\,(g_k+\log p_k)=m\right)=p_m\]

First, let \(\phi^{(m)}\) be the logit corresponding to token \(m\), so that

\[p_m=\Pr(X=m)=\frac{e^{\phi^{(m)}}}{\sum_i e^{\phi^{(i)}}} =\operatorname{softmax}(\phi)_m, \qquad p_m\propto e^{\phi^{(m)}}\]

Set

\[Y(m)=\phi^{(m)}+g_m\]

We want to show that

\[\Pr\!\left(\underset{k}{\arg\max}\,[\phi^{(k)}+g_k]=m\right) =\operatorname{softmax}(\phi)_m\]

\(Y(m)\) follows a shifted Gumbel distribution, with density

\[f_m(x)=e^{\phi^{(m)}-x}\exp\!\left(-e^{\phi^{(m)}-x}\right)\]

Then

\[ \begin{aligned} \Pr\!\left(\underset{k}{\arg\max}\,Y(k)=m\right) &=\Pr\!\left(Y(m)>Y(k),\ \forall k\ne m\right)\\ &=\int_{-\infty}^{\infty} \Pr\!\left(\bigcap_{k\ne m}\{\phi^{(k)}+g_k\le x\}\right)f_m(x)\,dx\\ &=\int_{-\infty}^{\infty} \prod_{k\ne m}\Pr\!\left(g_k\le x-\phi^{(k)}\right)f_m(x)\,dx\\ &=\int_{-\infty}^{\infty} \left(\prod_{k\ne m}e^{-e^{-x+\phi^{(k)}}}\right) e^{\phi^{(m)}-x}e^{-e^{-x+\phi^{(m)}}}\,dx\\ &=\int_{-\infty}^{\infty} e^{\phi^{(m)}-x} \exp\!\left(-\sum_k e^{-x+\phi^{(k)}}\right)\,dx\\ &=\int_{-\infty}^{\infty} e^{\phi^{(m)}}e^{-x} \exp\!\left(-e^{-x}\sum_k e^{\phi^{(k)}}\right)\,dx.\end{aligned}\]

With \(u=e^{-x}\sum_k e^{\phi^{(k)}}\), the last integral becomes

\[ \begin{aligned} \Pr\!\left(\underset{k}{\arg\max}\,Y(k)=m\right) &=\frac{e^{\phi^{(m)}}}{\sum_k e^{\phi^{(k)}}} \int_0^\infty e^{-u}\,du\\ &=\frac{e^{\phi^{(m)}}}{\sum_k e^{\phi^{(k)}}}\\ &=\operatorname{softmax}(\phi)_m. \qquad\square\end{aligned}\]

We just proved that

\[ \Pr\!\left(\underset{k}{\arg\max}\,[g_k+\phi^{(k)}]=m\right) =\operatorname{softmax}(\phi)_m =\Pr(X=m)\]

This means that

\[ \Pr\!\left(\underset{k}{\arg\max}\,[g_k+\log p_k]=m\right) =\operatorname{softmax}(\phi)_m\]

because

\[ \log p_k=\phi^{(k)}+\underbrace{\log\!\left(\frac{1}{\sum_i e^{\phi^{(i)}}}\right)}_{\text{constant in }k}\]

Now, to sample a new token, instead of using softmax followed by categorical sampling, we can simply add Gumbel noise to the logits and take the argmax. We obtain the same token distribution.

Now observe that this Gumbel-Max trick is equivalent to the following. Sample

\[R=(R_1,\ldots,R_K),\qquad R_k\overset{\mathrm{iid}}{\sim}\operatorname{Uniform}(0,1)\]

and set

\[m=\underset{k}{\arg\max}\;R_k^{1/p_k}\]

Indeed, using \(g_k=-\log(-\log R_k)\),

\[\begin{aligned} \underset{k}{\arg\max}\,[\log p_k+g_k] &=\underset{k}{\arg\max}\,[\log p_k-\log(-\log R_k)]\\ &=\underset{k}{\arg\max}\;\log\!\left(\frac{p_k}{-\log R_k}\right)\\ &=\underset{k}{\arg\max}\;\frac{p_k}{-\log R_k}\\ &=\underset{k}{\arg\min}\;\frac{-\log R_k}{p_k}\\ &=\underset{k}{\arg\max}\;\frac{\log R_k}{p_k}\\ &=\underset{k}{\arg\max}\;\log\!\left(R_k^{1/p_k}\right)\\ &=\underset{k}{\arg\max}\;R_k^{1/p_k} \end{aligned}\]

We will continue with this second formulation for the rest of this section.

The Inverse Transform Method

Provided that computers can give us independent and identically distributed random variables that follow \(U\sim\mathcal{U}(0,1)\), we need to find a way to sample from any distribution, including the token distribution, which turns out to be categorical. Top-k/top-p sampling uses the exponential distribution, so we also need to know how to sample from it. This is basically possible thanks to the inverse transform method.

Let \(X\) be a random variable with cumulative distribution function (CFD) \(F(x)\). The function \(F\) is defined on \(\mathbb{R}\rightarrow[0,1]\) as \(F(x)=P(X\leq x)\), and is thus a non-decreasing function.

The generalized inverse CDF is defined as:

\(F^{-1}(y)=\inf\{x\in\mathbb{R}:F(x)\geq y\},\quad y\in(0,1)\)

Let \(X=F^{-1}(U)\), where \(U\sim\mathcal{U}(0,1)\). Then \(X\) has CDF \(F\).

In fact, when \(F\) is invertible, we have:

\(P(X\leq x)=P(F^{-1}(U)\leq x)=P(U\leq F(x))=F(x)\)

Thus, to generate a random variable \(X\) with CDF \(F\), we can sample \(U\sim\mathcal{U}(0,1)\) and set \(X=F^{-1}(U)\).

Sampling from a Categorical Distribution

For instance, if we consider the token distribution \(p=(p_1,\ldots,p_K)\), where:

\(P(X=i)=p_i,\quad \forall i=1,\ldots,K,\quad \sum_{i=1}^{K}p_i=1\)

Its cumulative distribution function is:

\(F(x)=P(X\leq x)=\sum_{i\leq x}p_i\)

We can sample \(U\sim\mathcal{U}(0,1)\) using a pseudorandom generator.

Then, if \(U\leq p_1\), we return \(X=1\). Otherwise, we return \(X=i\) if:

\(\sum_{j=1}^{i-1}p_j<U\leq\sum_{j=1}^{i}p_j\)

That is, we find the smallest positive integer \(k\) such that \(U\leq F(k)\) and return \(X=k\).

Sampling from an Exponential Distribution

Now consider an exponential distribution with rate \(\lambda>0\), whose CDF is:

\(F(x)=1-e^{-\lambda x},\quad x\geq 0\)

Its inverse CDF is:

\(F^{-1}(y)=-\frac{1}{\lambda}\ln(1-y)\)

Thus, we can sample \(U\sim\mathcal{U}(0,1)\) using a pseudorandom generator and set:

\(X=-\frac{\ln(1-U)}{\lambda}\)

Alternatively, since \(1-U\sim\mathcal{U}(0,1)\), we can also write:

\(X=-\frac{\ln(U)}{\lambda}\)

The Exponential Race

Let \(X\) follow an exponential distribution with rate \(\lambda\), i.e., \(X\sim\operatorname{Exp}(\lambda)\).

Its probability density function is:

\(f_X(x)=\lambda e^{-\lambda x},\quad x\geq 0\)

And its cumulative distribution function is:

\(F_X(x)=1-e^{-\lambda x},\quad x\geq 0\)

If \(X_1,\ldots,X_n\) are independent and identically distributed exponential random variables with rate \(\lambda\), then:

\(S_n=X_1+\cdots+X_n\sim\operatorname{Gamma}(n,\lambda)\)

Additionally, if \(X\sim\operatorname{Exp}(1/\lambda)\), then:

\(X/\lambda\sim\operatorname{Exp}(1)\)

Now, in our context, let \(E_i\sim\operatorname{Exp}(1)\) be independent exponential random variables.

Then:

\(T_i=\frac{E_i}{p_i}\sim\operatorname{Exp}(p_i)\)

Now let's look at:

\(k=\operatorname*{arg\,max}_i\frac{p_i}{E_i}\)

Taking the logarithm, which preserves the ordering:

\(k=\operatorname*{arg\,max}_i\log\left(\frac{p_i}{E_i}\right)=\operatorname*{arg\,max}_i\left(\log p_i-\log E_i\right)\)

But since \(E_i\sim\operatorname{Exp}(1)\), we have:

\(G_i=-\log E_i\sim\operatorname{Gumbel}(0,1)\)

Therefore:

\(k=\operatorname*{arg\,max}_i\left(\log p_i+G_i\right)\)

This is exactly the Gumbel-Max trick, since \(P\left(k=\operatorname*{argmin}_i T_i\right)=p_k\), assuming that \(\sum_{i=1}^{K}p_i=1\). Thus, this provides another way to sample tokens from a categorical distribution without distorting the original token probabilities!.

Gumbel-Max Watermarking

Now we have a new trick to sample a token. Given \(p=(p_1,\ldots,p_K)\), let

\[ R=(R_1,\ldots,R_K),\qquad R_v\overset{\mathrm{iid}}{\sim}\operatorname{Uniform}(0,1) \]

and pick

\[ V^*=\underset{v}{\arg\max}\;R_v^{1/p_v}\]

\(V^*\) is a random variable over the set of tokens. We have that, conditionally on \(V^*=v\),

\[ R_{V^*}\mid(V^*=v)\sim\operatorname{Beta}(1/p_v,1)\]

Let us prove it.

For \(R_v\sim\operatorname{Uniform}(0,1)\), let \(Y_v=R_v^{1/p_v}\). Then

\[ \begin{aligned} \Pr(Y_v\le t) &=\Pr\!\left(R_v^{1/p_v}\le t\right)\\ &=\Pr(R_v\le t^{p_v})\\ &=t^{p_v} \end{aligned}\]

Therefore,

\[ Y_v\sim\operatorname{Beta}(p_v,1)\]

Now, \[ Y_{V^*}=\max_v Y_v\sim\operatorname{Uniform}(0,1)\]

since the maximum of these independent Beta variables has CDF

\[ \begin{aligned} \Pr(Y_{V^*}\le t) &=\prod_v\Pr(Y_v\le t)\\ &=\prod_v t^{p_v}\\ &=t^{\sum_vp_v}=t \end{aligned} \]

Moreover:

\[ \begin{aligned} \Pr(Y_{V^*}\le t,\ V^*=v) &= \int_0^t P(Y_j \leq y,\ \forall j \neq v \mid Y_v = y)\,f_{Y_v}(y)\,dy &=\int_0^t p_v y^{p_v-1}\prod_{j\ne v}y^{p_j}\,dy\\ &=p_v\int_0^t y^{\sum_jp_j-1}\,dy\\ &=p_vt,\\ \Pr(Y_{V^*}\le t\mid V^*=v) &=\frac{p_vt}{\Pr(V^*=v)}\\ &=t, \end{aligned} \]

where \(\Pr(V^*=v)=p_v\) by the Gumbel-Max trick.

Now \(R_{V^*}=(Y_{V^*})^{p_{V^*}}\). Conditionally on \(V^*=v\),

\[ \begin{aligned} \Pr(R_{V^*}\le t\mid V^*=v) &=\Pr\!\left((Y_{V^*})^{p_v}\le t\mid V^*=v\right)\\ &=\Pr\!\left(Y_{V^*}\le t^{1/p_v}\mid V^*=v\right)\\ &=t^{1/p_v} \end{aligned}\]

Thus,

\[ R_{V^*}\mid(V^*=v)\sim\operatorname{Beta}(1/p_v,1).\qquad\square\]

We have just proved that \(R_{V^*}\) follows \(\operatorname{Beta}(1/p_v,1)\) conditionally on \(V^*=v\). This is key for detection.

Let us define the score

\[ S=\sum_{t=1}^T-\ln(1-R_t)\]

Instead of using freshly sampled independent \(R_v\) values, we use

\[ R_v=\operatorname{PRF}(v,w,K)\]

where \(w\) is the context window consisting of the last \(\gamma\) tokens and \(K\) is the secret key. The PRF supplies pseudorandom values in \([0,1]\).

At time step \(t\), the token is selected by

\[ x_t=\underset{v\in\mathcal V}{\arg\max}\; \left(R_v^{(t)}\right)^{1/p_v^{(t)}}\]

Now, given an input token stream \(x^{(1)},\ldots,x^{(T)}\), we compute the PRF values \(R_v^{(t)}\in[0,1]\) using the context at each position \(t\). We then select the PRF value corresponding to the observed token. Define

\[R_t:=R_{x^{(t)}}^{(t)}\]

that is, the PRF value of the observed token at position \(t\). The detection score is

\[S_T=-\sum_{t=1}^T\ln(1-R_t)\]

Watermarked tokens tend to have high \(R_t\) values, so \(S_T\) tends to be high. For non-watermarked text the \(R_t\) values are uniform. We can therefore use a detection threshold \(\tau\) and decide that the text is watermarked if

\[ S_T>\tau\]

In fact, under \(H_0\),

\[-\ln(1-R_t)\sim\operatorname{Exp}(1)\]

so, assuming independent scores:

\[S_T\mid H_0\sim\operatorname{Gamma}(T,1),\]

where the second parameter denotes the rate. Therefore,

\[\mathbb E[S_T\mid H_0]=T\]

Under the watermark hypothesis \(H_1\), the expectation is larger. In particular, for token probabilities \(p_v^{(t)}\), the entropy is

\[H_t=-\sum_v p_v^{(t)}\ln p_v^{(t)}\]

and a lower bound on the expected score is [6]

\[\mathbb E[S_T\mid H_1]\ge T+\left(\frac{\pi^2}{6}-1\right)\sum_{t=1}^T H_t\]

Thus, when the distribution has low entropy (the model is confident), the watermark signal is weak. While on the other hand, when the LLM is uncertain (i.e high entropy), the watermark produces a higher signal.

Instead of using a direct threshold on the score, we can express detection through a p-value:

\[\operatorname{p\text{-}value}(s) =\Pr(S_T\ge s\mid H_0) =\frac{\Gamma(T,s)}{\Gamma(T)}\]

Here, \(\Gamma(T,s)\) is the upper incomplete gamma function; this allows us to control the false-positive rate by comparing the p-value with a chosen significance level \(\alpha\).

vLLM implementation

Let's look at the vLLM implementation [6].

Gumbel watermarking is implemented by GumbelWatermarker.

class GumbelWatermarker(Watermarker):
    def __init__(
        self,
        key: int,
        context_width: int = 4,
        prf: WatermarkPRF | WatermarkPRFName = "philox",
    ) -> None:
        prf = create_prf(prf, key) if isinstance(prf, str) else prf
        _validate_context_width(context_width)
        self.prf = prf
        self._context_width = context_width

The sampling procedure follows exactly what we covered, as shown below. Uniform values are first generated using the PRF, uniforms = self.prf.uniform(contexts, vocabulary), for each token in the vocabulary. Then, Gumbel noise is constructed using noise = -torch.log(-torch.log(uniforms)), followed by token sampling using torch.argmax(logits + noise, dim=-1). The call self.prf.uniform(contexts, vocabulary) is deterministic for the same context, assuming the key remains unchanged.

def _sample_watermarked(
    self,
    logits: torch.Tensor,
    contexts: torch.Tensor,
) -> WatermarkSample:
    if type(self.prf) is PhiloxPRF and logits.device.type == "cuda":
        from vllm.v1.worker.gpu.sample.watermark import philox_gumbel_sample
        return WatermarkSample(
            philox_gumbel_sample(logits, contexts, self.prf.key), logits
        )

    vocabulary = torch.arange(logits.shape[-1], device=logits.device)
    uniforms = self.prf.uniform(contexts, vocabulary)
    uniforms = uniforms.clamp_min(torch.finfo(torch.float32).tiny)
    noise = -torch.log(-torch.log(uniforms))
    return WatermarkSample(torch.argmax(logits + noise, dim=-1), logits)

During detection, the contexts are reconstructed, and a score is computed for each target token.

def _score_tokens(
    self, contexts: torch.Tensor, targets: torch.Tensor
) -> torch.Tensor:
    uniforms = self.prf.uniform(contexts, targets.unsqueeze(-1)).squeeze(-1)
    return -torch.log1p(-uniforms.to(torch.float64))

Next, the p-value is obtained using the Gamma function.

def detect(self, token_ids: list[int]) -> WatermarkDetection:
    contexts, targets = self._prepare_inputs(token_ids)
    num_scored_tokens = len(targets)
    if num_scored_tokens == 0:
        return WatermarkDetection(0.0, 1.0, 0, False)

    token_scores = self._score_tokens(contexts, targets)
    score = self._aggregate_scores(token_scores)
    p_value = self._get_p_value(score, num_scored_tokens)
    return WatermarkDetection(
        score=score,
        p_value=p_value,
        num_scored_tokens=num_scored_tokens,
        is_watermarked=p_value <= self.p_value_threshold,
    )

When speculative decoding is enabled, i.e., when a draft model and rejection sampling are used, two watermark keys are employed, one for the draft model and another for the target model (and also when draft tokens are rejected). This is done to preserve the acceptance rate [7, 9].

If we look at the vLLM implementation, two watermark keys are used.

def create_speculative_draft_watermarker(
    watermarker: Watermarker,
    max_num_reqs: int,
    device: torch.device,
    allow_target_only: bool,
    num_speculative_steps: int = 1,
    deduplicate_contexts: WatermarkContextScope = "none",
    deduplicate_contexts_max_history: int | None = 8192,
) -> DraftWatermarker | None:
    if isinstance(watermarker, SupportsSpeculativeDecoding):
        return DraftWatermarker(
            watermarker.draft_watermarker,
            max_num_reqs,
            device,
            num_speculative_steps,
            deduplicate_contexts,
            deduplicate_contexts_max_history,
        )
    if allow_target_only:
        return None
    raise ValueError(
        f"The {type(watermarker).__name__} watermarking algorithm does not support "
        "speculative decoding. Set allow_target_only_watermarking=true to "
        "leave draft tokens unwatermarked."
    )

class DualKeyGumbelWatermarker(Watermarker, SupportsSpeculativeDecoding):
    def __init__(
        self,
        key: int,
        context_width: int = 4,
        prf: WatermarkPRFName = "philox",
        alpha: float = 0.1,
    ) -> None:
        if not 0 <= alpha <= 1:
            raise ValueError("alpha must be between 0 and 1")
        self.alpha = alpha
        self.draft_watermarker = GumbelWatermarker(
            derive_watermark_key(key, b"key_a"), context_width, prf
        )
        self.target_watermarker = GumbelWatermarker(
            derive_watermark_key(key, b"key_b"), context_width, prf
        )
        self._routing_logits_cache: dict[torch.device, torch.Tensor] = {}

During draft sampling, the logits and contexts are passed to the Gumbel sampler using the watermark key associated with the draft model.

class DraftWatermarker:
    def sample(
        self,
        logits: torch.Tensor,
        idx_mapping: torch.Tensor,
        temperature: torch.Tensor,
        seed: torch.Tensor,
        pos: torch.Tensor,
        apply_temperature: bool,
        is_drafting: bool,
        logits_cache: torch.Tensor | None = None,
        logits_cache_col: torch.Tensor | None = None,
        use_fp64: bool = False,
        logits_cache_source: torch.Tensor | None = None,
    ) -> torch.Tensor:
        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]
        steps, enabled = self._sampling_state(
            logits, idx_mapping, request_temperatures, logits_cache_col, contexts
        )
        processed_logits = logits / torch.where(
            request_temperatures == 0, 1, request_temperatures
        ).unsqueeze(-1)
        random_sampler = 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=steps,
            logits_cache_source=(
                logits if logits_cache_source is None else logits_cache_source
            ),
        )
        sampled = self.watermarker.sample(
            processed_logits,
            contexts,
            random_sampler=random_sampler,
            skip_mask=~enabled,
        ).token_ids
        contexts.copy_(torch.cat((contexts[:, 1:], sampled.unsqueeze(-1)), dim=-1))
        return sampled

On the other hand, the rejection sampler uses the target model's watermark key, that is, when a token is rejected the correction token is sample using the target watermark key; the same applies when all draft tokens are accepted and a bonus token is to be sampled.

self.rejection_sampler = RejectionSampler(
    self.sampler,
    self.speculative_config,
    self.device,
    watermark_key=speculative_target_watermark_key(
        self.vllm_config.watermark_config
    ),
)

...

def create_speculative_target_watermarker(watermarker):
    if isinstance(watermarker, SupportsSpeculativeDecoding):
        return watermarker.target_watermarker
    return watermarker

During detection, all tokens must be evaluated against both watermark keys (i.e scores must be computed under both keys), since there is no signal indicating which key was used for each token; naturally this reduces detection performance. If we look at vLLM, it aggregates them as a weighted sum.

def _score_tokens(
    self, contexts: torch.Tensor, targets: torch.Tensor
) -> torch.Tensor:
    scores = torch.stack(
        [
            -torch.log1p(
                -prf.uniform(contexts, targets.unsqueeze(-1))
                .squeeze(-1)
                .to(torch.float64)
            )
            for prf in (self.prf, self.key_b_prf)
        ],
        dim=-1,
    )
    weights = scores.new_tensor([1 - self.alpha, self.alpha])
    return scores @ weights

References and readings

[1] https://simons.berkeley.edu/sites/default/files/2024-10/LLM24-2%20Slides%20-%20Scott%20Aaronson.pdf

[2] https://github.com/vllm-project/vllm/issues/53916

[3] https://arxiv.org/pdf/2308.00113

[4] https://papers.nips.cc/paper_files/paper/2024/file/2567c95fd41459a98a73ba893775d22a-Paper-Conference.pdf

[5] https://www.research-collection.ethz.ch/server/api/core/bitstreams/1adfb599-4160-409a-ad27-dc0d4109abdc/content

[6] https://arxiv.org/pdf/2605.12456

[7] https://static-content.springer.com/esm/art%3A10.1038%2Fs41586-024-08025-4/MediaObjects/41586_2024_8025_MOESM1_ESM.pdf

[8] https://vllm.ai/blog/2026-09-24-watermarking-in-vllm

[9] https://github.com/vllm-project/vllm/pull/56122

[10] https://www.nature.com/articles/s41586-024-08025-4.pdf