vllm.v1.watermarking.spec_decode

Functions:

• speculative_target_watermark_key –

  Resolve the resample kernel's Philox key from configuration.

• watermarked_rejection_sample –

  Use the target key except at contexts selected for ordinary sampling.

_resolve_watermark_key(watermarker)

Return the Philox recovery key.

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

def _resolve_watermark_key(watermarker: Watermarker) -> int:
    """Return the Philox recovery key."""
    if isinstance(watermarker, SupportsSpeculativeDecoding):
        raise ValueError(
            f"{type(watermarker).__name__} keys the target role separately from "
            "the draft. Pass create_speculative_target_watermarker(watermarker) "
            "so the recovery draw carries the target's key."
        )
    return _philox_key(watermarker)

speculative_target_watermark_key(watermark_config)

Resolve the resample kernel's Philox key from configuration.

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

def speculative_target_watermark_key(
    watermark_config: WatermarkConfig | None,
) -> int | None:
    """Resolve the resample kernel's Philox key from configuration."""
    if watermark_config is None:
        return None

    return _resolve_watermark_key(
        create_speculative_target_watermarker(create_watermarker(watermark_config))
    )

watermarked_rejection_sample(target_logits, draft_logits, draft_sampled, cu_num_logits, pos, idx_mapping, expanded_idx_mapping, expanded_local_pos, temperature, seed, num_speculative_steps, contexts, watermarking, watermarker, watermarking_skip_mask=None, use_fp64=False)

Use the target key except at contexts selected for ordinary sampling.

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

def watermarked_rejection_sample(
    target_logits: torch.Tensor,
    draft_logits: torch.Tensor | None,
    draft_sampled: torch.Tensor,
    cu_num_logits: torch.Tensor,
    pos: torch.Tensor,
    idx_mapping: torch.Tensor,
    expanded_idx_mapping: torch.Tensor,
    expanded_local_pos: torch.Tensor,
    temperature: torch.Tensor,
    seed: torch.Tensor,
    num_speculative_steps: int,
    contexts: torch.Tensor,
    watermarking: torch.Tensor,
    watermarker: Watermarker,
    watermarking_skip_mask: torch.Tensor | None = None,
    use_fp64: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Use the target key except at contexts selected for ordinary sampling."""
    assert contexts.shape == (target_logits.shape[0], watermarker.context_width)
    return rejection_sample(
        target_logits,
        draft_logits,
        draft_sampled,
        cu_num_logits,
        pos,
        idx_mapping,
        expanded_idx_mapping,
        expanded_local_pos,
        temperature,
        seed,
        num_speculative_steps,
        use_fp64=use_fp64,
        contexts=contexts,
        watermarking=watermarking,
        watermarking_skip_mask=watermarking_skip_mask,
        watermark_key=_resolve_watermark_key(watermarker),
    )