vLLM/vllm/v1/watermarking/gpu_sampler.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

from typing import Literal

import numpy as np
import torch

from vllm.logger import init_logger
from vllm.sampling_params import SamplingParams
from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p
from vllm.v1.watermarking.watermarker import RandomSampler, Watermarker
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 repeated_context_mask

logger = init_logger(__name__)


class GPUWatermarkSampler(Sampler):
    def __init__(
        self,
        watermarker: Watermarker,
        *args,
        deduplicate_contexts: Literal["none", "single_turn", "all"] = "single_turn",
        deduplicate_contexts_max_history: int | None = 8192,
        **kwargs,
    ) -> None:
        super().__init__(*args, **kwargs)
        self.watermarker = watermarker
        self.deduplicate_contexts = deduplicate_contexts
        self.deduplicate_contexts_max_history = deduplicate_contexts_max_history
        self.watermarking = UvaBackedTensor(
            self.sampling_states.max_num_reqs, dtype=torch.bool
        )
        self.watermarking.np.fill(True)
        self.watermarking.copy_to_uva()

    def add_request(self, req_idx: int, sampling_params: SamplingParams) -> None:
        super().add_request(req_idx, sampling_params)
        self.watermarking.np[req_idx] = sampling_params.watermarking
        if sampling_params.watermarking and sampling_params.temperature == 0:
            logger.warning_once(
                "Watermarking is enabled, but greedy decoding "
                "(temperature=0) cannot be watermarked. This request will use "
                "ordinary greedy sampling."
            )

    def apply_staged_writes(self) -> None:
        super().apply_staged_writes()
        self.watermarking.copy_to_uva()

    def _sample_random(
        self,
        processed_logits: torch.Tensor,
        expanded_idx_mapping: torch.Tensor,
        idx_mapping_np: np.ndarray,
        pos: torch.Tensor,
        top_k: torch.Tensor | None,
        top_p: torch.Tensor | None,
        use_fused_sampler: bool,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        enabled = self.watermarking.np[idx_mapping_np] & (
            self.sampling_states.temperature.np[idx_mapping_np] != 0
        )
        if not np.any(enabled):
            return super()._sample_random(
                processed_logits,
                expanded_idx_mapping,
                idx_mapping_np,
                pos,
                top_k,
                top_p,
                use_fused_sampler,
            )

        processed_logits = apply_top_k_top_p(processed_logits, top_k, top_p)
        contexts = self._get_contexts(expanded_idx_mapping)
        repeated_contexts = None
        if self.deduplicate_contexts != "none":
            repeated_contexts = self._get_repeated_contexts(
                expanded_idx_mapping, contexts
            )

        temperatures = self.sampling_states.temperature.gpu[expanded_idx_mapping]
        needs_mixed_sampling = repeated_contexts is not None or not np.all(enabled)
        skip_mask = None
        if needs_mixed_sampling:
            watermarking = self.watermarking.gpu[expanded_idx_mapping] & (
                temperatures != 0
            )
            if repeated_contexts is not None:
                watermarking &= ~repeated_contexts
            skip_mask = ~watermarking

        random_sampler = RandomSampler(
            expanded_idx_mapping=expanded_idx_mapping,
            temperatures=self.sampling_states.temperature.gpu,
            seeds=self.sampling_states.seeds.gpu,
            positions=pos,
            use_fp64=self.use_fp64_gumbel,
        )
        output = self.watermarker.sample(
            processed_logits,
            contexts,
            random_sampler,
            skip_mask=skip_mask,
        )
        sampled = output.token_ids
        output_logits = output.logits
        sampled = torch.where(
            temperatures == 0,
            processed_logits.argmax(dim=-1),
            sampled,
        )
        return sampled, output_logits

    def _get_repeated_contexts(
        self,
        expanded_idx_mapping: torch.Tensor,
        contexts: torch.Tensor,
    ) -> torch.Tensor:
        return repeated_context_mask(
            self.req_states.all_token_ids.gpu,
            expanded_idx_mapping,
            self.req_states.prompt_len.gpu,
            self.req_states.total_len.gpu,
            contexts,
            self.deduplicate_contexts_max_history,
            include_prompt=self.deduplicate_contexts == "all",
            skip_partial_context=self.deduplicate_contexts == "all",
        )

    def _get_contexts(
        self,
        expanded_idx_mapping: torch.Tensor,
        expanded_local_pos: torch.Tensor | None = None,
        draft_sampled: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """Build contexts for flattened verification rows.

        With committed tokens ``[10, 11]``, drafts ``[20, 21]``, and width 2,
        local positions ``[0, 1, 2]`` use contexts ``[10, 11]``, ``[11, 20]``,
        and ``[20, 21]``. The first verification row precedes the first draft.
        """
        context_width = self.watermarker.context_width
        req_indices = expanded_idx_mapping.to(torch.int64)
        valid_reqs = req_indices >= 0
        safe_req_indices = req_indices.clamp_min(0)
        total_lens = self.req_states.total_len.gpu[safe_req_indices].to(torch.int64)
        prompt_lens = self.req_states.prompt_len.gpu[safe_req_indices].to(torch.int64)
        if expanded_local_pos is None:
            expanded_local_pos = torch.zeros_like(req_indices)
        else:
            expanded_local_pos = expanded_local_pos.to(torch.int64)
        offsets = torch.arange(
            -context_width, 0, dtype=torch.int64, device=req_indices.device
        )
        positions = (
            total_lens.unsqueeze(-1) + expanded_local_pos.unsqueeze(-1) + offsets
        )
        committed = positions < total_lens.unsqueeze(-1)
        valid_committed = (
            valid_reqs.unsqueeze(-1)
            & committed
            & (positions >= prompt_lens.unsqueeze(-1))
        )
        committed_contexts = self.req_states.all_token_ids.gpu[
            safe_req_indices.unsqueeze(-1),
            positions.clamp(0, self.req_states.all_token_ids.gpu.shape[1] - 1),
        ]
        contexts = torch.where(valid_committed, committed_contexts, -1)
        if draft_sampled is None:
            return contexts

        row_indices = torch.arange(
            len(req_indices), dtype=torch.int64, device=req_indices.device
        )
        draft_offsets = positions - total_lens.unsqueeze(-1)
        draft_indices = (
            row_indices.unsqueeze(-1)
            - expanded_local_pos.unsqueeze(-1)
            + draft_offsets
            + 1
        )
        valid_drafts = (
            valid_reqs.unsqueeze(-1)
            & ~committed
            & (draft_offsets < expanded_local_pos.unsqueeze(-1))
        )
        draft_contexts = draft_sampled[draft_indices.clamp(0, len(draft_sampled) - 1)]
        return torch.where(valid_drafts, draft_contexts, contexts)