watermark_sample_warmup - vLLM
Skip to content

vllm.model_executor.warmup.watermark_sample_warmup

Warm up the watermarked sampler's Triton kernel.

_philox_gumbel_kernel is JIT-compiled once per (key, logits dtype, skip-mask, USE_FP64) specialization. The generic sampler warmup samples with SamplingParams.for_sampler_warmup(), whose feature-heavy logits processing forces the fp32 copy in apply_sampling_params, so the first ordinary temperature-1.0 request is the first launch with model-dtype logits and pays the compilation inside inference. This pre-compiles every specialization the sampler path can launch for the configured watermark.

_philox_key(watermarker)

None when the watermarker takes the torch path and compiles no kernel.

Source code in vllm/model_executor/warmup/watermark_sample_warmup.py 
32
33
34
35
36
37
38
39
40
def _philox_key(watermarker: Watermarker) -> int | None:
    """None when the watermarker takes the torch path and compiles no kernel."""
    from vllm.v1.watermarking.gumbel import GumbelWatermarker
    from vllm.v1.watermarking.prfs import PhiloxPRF

    if not isinstance(watermarker, GumbelWatermarker):
        return None
    prf = watermarker.prf
    return prf.key if type(prf) is PhiloxPRF else None

_philox_sampler_keys(watermarker)

Philox keys the sampler launches the kernel with.

Source code in vllm/model_executor/warmup/watermark_sample_warmup.py 
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
def _philox_sampler_keys(watermarker: Watermarker) -> list[int]:
    """Philox keys the sampler launches the kernel with."""
    from vllm.v1.watermarking.gumbel import DualKeyGumbelWatermarker

    watermarkers: list[Watermarker] = [watermarker]
    if isinstance(watermarker, DualKeyGumbelWatermarker):
        # ``DualKeyGumbelWatermarker.sample`` launches once per key, and
        # collapses to a single key at the alpha bounds, where the other key's
        # specialization is never launched.
        if watermarker.alpha == 0:
            watermarkers = [watermarker.draft_watermarker]
        elif watermarker.alpha == 1:
            watermarkers = [watermarker.target_watermarker]
        else:
            watermarkers = [
                watermarker.draft_watermarker,
                watermarker.target_watermarker,
            ]
    return [key for key in map(_philox_key, watermarkers) if key is not None]