@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 155 lines 6.35 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Warm up the watermarked sampler's Triton kernel.
4
5 ``_philox_gumbel_kernel`` is JIT-compiled once per (key, logits dtype,
6 skip-mask, ``USE_FP64``) specialization. The generic sampler warmup samples with
7 ``SamplingParams.for_sampler_warmup()``, whose feature-heavy logits processing
8 forces the fp32 copy in ``apply_sampling_params``, so the first ordinary
9 temperature-1.0 request is the first launch with model-dtype logits and pays
10 the compilation inside inference. This pre-compiles every specialization the
11 sampler path can launch for the configured watermark.
12 """
13
14 from __future__ import annotations
15
16 from typing import TYPE_CHECKING, cast
17
18 import torch
19
20 from vllm.logger import init_logger
21 from vllm.platforms import current_platform
22
23 if TYPE_CHECKING:
24 from vllm.v1.watermarking.watermarker import Watermarker
25 from vllm.v1.worker.gpu_worker import Worker
26
27 logger = init_logger(__name__)
28
29 _NUM_WARMUP_TOKENS = 1
30
31
32 def _philox_key(watermarker: Watermarker) -> int | None:
33 """None when the watermarker takes the torch path and compiles no kernel."""
34 from vllm.v1.watermarking.gumbel import GumbelWatermarker
35 from vllm.v1.watermarking.prfs import PhiloxPRF
36
37 if not isinstance(watermarker, GumbelWatermarker):
38 return None
39 prf = watermarker.prf
40 return prf.key if type(prf) is PhiloxPRF else None
41
42
43 def _philox_sampler_keys(watermarker: Watermarker) -> list[int]:
44 """Philox keys the sampler launches the kernel with."""
45 from vllm.v1.watermarking.gumbel import DualKeyGumbelWatermarker
46
47 watermarkers: list[Watermarker] = [watermarker]
48 if isinstance(watermarker, DualKeyGumbelWatermarker):
49 # ``DualKeyGumbelWatermarker.sample`` launches once per key, and
50 # collapses to a single key at the alpha bounds, where the other key's
51 # specialization is never launched.
52 if watermarker.alpha == 0:
53 watermarkers = [watermarker.draft_watermarker]
54 elif watermarker.alpha == 1:
55 watermarkers = [watermarker.target_watermarker]
56 else:
57 watermarkers = [
58 watermarker.draft_watermarker,
59 watermarker.target_watermarker,
60 ]
61 return [key for key in map(_philox_key, watermarkers) if key is not None]
62
63
64 @torch.inference_mode()
65 def watermark_sample_warmup(worker: Worker) -> None:
66 if getattr(worker.vllm_config, "watermark_config", None) is None:
67 return
68 # ``GumbelWatermarker`` launches the kernel only on CUDA
69 # (vllm/v1/watermarking/gumbel.py:63,:82); elsewhere the sampler takes the
70 # torch path and warming here would compile a kernel nothing launches.
71 if not current_platform.is_cuda_alike():
72 return
73
74 from vllm.v1.watermarking.gpu_sampler import GPUWatermarkSampler
75 from vllm.v1.worker.gpu.sample.watermark import philox_gumbel_sample
76
77 # Read the sampler the runner built instead of rebuilding it from config:
78 # its watermarker (the target role under a speculative config), its context
79 # deduplication setting and its fp64 flag are what the launches have to
80 # match, and cannot drift from the runner. A rank that does not sample -- a
81 # non-final pipeline stage, a pooling model -- holds no such sampler.
82 sampler = getattr(worker.model_runner, "sampler", None)
83 if not isinstance(sampler, GPUWatermarkSampler):
84 return
85
86 model_config = worker.vllm_config.model_config
87 device = worker.device
88 try:
89 keys = _philox_sampler_keys(sampler.watermarker)
90 if not keys:
91 return
92
93 # The sampler passes model-dtype logits through when no logits
94 # processing is needed and fp32 otherwise. ``ModelConfig.dtype`` is
95 # resolved to a torch dtype during validation.
96 dtypes = {cast("torch.dtype", model_config.dtype), torch.float32}
97 # ``GPUWatermarkSampler._sample_random`` builds a skip mask whenever
98 # context deduplication is on; with it off only a batch that mixes in an
99 # opted-out or greedy row does, and the mask-free specialization stays
100 # reachable.
101 with_skip_mask = (
102 (True,) if sampler.deduplicate_contexts != "none" else (True, False)
103 )
104
105 vocab_size = model_config.get_vocab_size()
106 contexts = torch.zeros(
107 (_NUM_WARMUP_TOKENS, sampler.watermarker.context_width),
108 dtype=torch.int32,
109 device=device,
110 )
111 # Dtypes mirror the runtime buffers the sampler passes and are part of
112 # the specialization key: int32 request-state token ids, int64 index and
113 # seed buffers, fp32 sampling temperatures.
114 sampling_state = {
115 "skip_mask": torch.zeros(
116 _NUM_WARMUP_TOKENS, dtype=torch.bool, device=device
117 ),
118 "expanded_idx_mapping": torch.zeros(
119 _NUM_WARMUP_TOKENS, dtype=torch.int64, device=device
120 ),
121 "temperatures": torch.ones(
122 _NUM_WARMUP_TOKENS, dtype=torch.float32, device=device
123 ),
124 "seeds": torch.zeros(_NUM_WARMUP_TOKENS, dtype=torch.int64, device=device),
125 "positions": torch.zeros(
126 _NUM_WARMUP_TOKENS, dtype=torch.int64, device=device
127 ),
128 }
129
130 logger.info(
131 "Warming up watermark sampler kernel (vocab=%d, keys=%d, dtypes=%s, "
132 "skip_mask=%s).",
133 vocab_size,
134 len(keys),
135 [str(dtype) for dtype in dtypes],
136 list(with_skip_mask),
137 )
138 for dtype in dtypes:
139 logits = torch.zeros(
140 (_NUM_WARMUP_TOKENS, vocab_size), dtype=dtype, device=device
141 )
142 for key in keys:
143 for use_skip_mask in with_skip_mask:
144 philox_gumbel_sample(
145 logits,
146 contexts,
147 key,
148 # Only the mixed path forwards the sampler's setting;
149 # ``GumbelWatermarker._sample_watermarked`` launches
150 # the mask-free variant with the fp32 default.
151 use_fp64=use_skip_mask and sampler.use_fp64_gumbel,
152 **(sampling_state if use_skip_mask else {}),
153 )
154 except Exception:
155 logger.warning("Skipping watermark sampler warmup.", exc_info=True)