vllm.v1.watermarking.watermarker

Classes:

• SupportsSpeculativeDecoding –

  A watermarker compatible with standard speculative rejection sampling.

• Watermarker –

SupportsSpeculativeDecoding

Bases: Protocol

A watermarker compatible with standard speculative rejection sampling.

Source code in vllm/v1/watermarking/watermarker.py

@runtime_checkable
class SupportsSpeculativeDecoding(Protocol):
    """A watermarker compatible with standard speculative rejection sampling."""

    draft_watermarker: Watermarker
    target_watermarker: Watermarker

Watermarker

Bases: ABC

Source code in vllm/v1/watermarking/watermarker.py

class Watermarker(ABC):
    @property
    @abstractmethod
    def context_width(self) -> int:
        raise NotImplementedError

    def sample(
        self,
        logits: torch.Tensor,
        contexts: torch.Tensor,
        random_sampler: RandomSampler | None = None,
        skip_mask: torch.Tensor | None = None,
    ) -> WatermarkSample:
        if skip_mask is not None:
            if random_sampler is None:
                raise ValueError("skip_mask requires a random sampler")
            mixed = self._try_sample_mixed(logits, contexts, skip_mask, random_sampler)
            if mixed is not None:
                return mixed

        watermarked = self._sample_watermarked(logits, contexts)
        if skip_mask is None:
            return watermarked

        assert random_sampler is not None
        token_ids = torch.where(
            skip_mask, random_sampler(logits), watermarked.token_ids
        )
        output_logits = watermarked.logits
        if output_logits is not logits:
            output_logits = torch.where(skip_mask.unsqueeze(-1), logits, output_logits)
        return WatermarkSample(token_ids, output_logits)

    @abstractmethod
    def _sample_watermarked(
        self,
        logits: torch.Tensor,
        contexts: torch.Tensor,
    ) -> WatermarkSample:
        raise NotImplementedError

    def _try_sample_mixed(
        self,
        logits: torch.Tensor,
        contexts: torch.Tensor,
        skip_mask: torch.Tensor,
        random_sampler: RandomSampler,
    ) -> 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

_try_sample_block(logits, num_steps, random_sampler, state)

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.

Source code in vllm/v1/watermarking/watermarker.py

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