vllm.v1.watermarking.spec_decode ¶

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)

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))
    )

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,
    use_fp64: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Use the target key for rejection recovery and bonus tokens."""
    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,
        watermark_key=_resolve_watermark_key(watermarker),
    )