Prhub

#56122 [watermarking] Dual-key gumbel-max watermarking for speculative decoding support

原始 PR 作者 TQCB 合并时间 2026-09-13 02:42 文件变更 28 提交数 35 评论 22 代码增减 +3451 / -236

执行摘要

水印兼容投机解码:双 key Gumbel 与上下文去重

PR body 明确指出现有水印(前置 PR #54053 的 Gumbel-max 单 key 方案)与投机解码不兼容,因此提出两条路径:allow_target_only_watermarking=true 允许草稿 token 不被水印(但会按接受比例稀释水印信号),以及新的 dual_key_gumbel 算法,让草稿与目标各自使用独立 key 水印,同时保持接受率不变(引用 SynthID-Text 补充材料 algorithm 5/6 与 K.13 节)。第二条动机来自上下文去重:默认 4 token 的短上下文在同一序列内可能重复,导致同一位置反复受到相同水印偏置,破坏“单序列非失真”性质,PR body 给出的实测是未去重时 35.7% 的生成撞上 1024 token 上限、4-token 上下文重复率 74.94%,而去重后分别降到 4.0% 与 17.99%。

值得精读,尤其是三类读者:做水印/内容溯源的人应关注 DualKeyGumbelWatermarker 的路由与检测器方差重标定;做投机解码的人应关注 resample_kernel 的 WATERMARK 分支与 watermarked_rejection_sample 的 key 角色约定;做性能的人应关注 context 去重的 kernel 代价与 watermark_sample_warmup 的 specialization 覆盖策略。建议重点看两个设计决策:一是 Watermarker.sample 模板方法 + _try_sample_mixed 的“可选融合钩子”,它把两趟采样优化收敛到算法内部;二是 create_speculative*_watermarker 的“拆角色”约定,配合 _resolve_watermark_key 的显式报错,避免双 key 容器被误用。阅读时可跳过 tests/watermarking 的数值断言细节,但 tests/v1/spec_decode/test_rejection_sampler_utils.py 中“水印恢复等价于 philox_gumbel_sample”的对比测试值得一看。

讨论亮点

1) 抽象层次之争(tomasruizt vs TQCB):tomasruizt 认为“it is not super easy to understand the relationship of different watermarking classes”,建议去掉 WatermarkedRejectionSampler 子类、给 DualKeyGumbelWatermarker 改用组合、把 DraftWatermarker 的 contexts/enabled 抽成状态对象;TQCB 回复接受前两条,但认为“DraftWatermarker handles quite a lot that we can't easily move to a state object, I think this will be difficult to flatten without a tonne of duplication”,并希望保留协议边界。最终 tomasruizt 提交的 25df1e46(删除 WatermarkedRejectionSampler,改为向 RejectionSampler 传可选 target watermark key)与 1153b1c0(DualKeyGumbelWatermarker 改用组合)被合入。
2) 27B 模型检测率低的归因(simon-veitner-redhat):复现后认为“the cause is how little entropy the 27B leaves in five-shot GSM8K answers, not the dual-key scheme, MTP, or the base feature”,并指出“On ordinary prose the 27B is detected in 99–100% of completions. Not a blocker; the description and docs should say this instead of guessing at it”,推动文档改写而不是改算法。
3) 内核缺陷与性能返工(simon-veitner-redhat):context mask kernel 末尾多读了 BLOCK-3 个位置,compute-sanitizer 报出 47 次非法全局读,“the comparison already excluded those lanes, so results were unaffected; the read was the defect”;key-B 恢复最初用“跑完 stock sampler 再用 PyTorch 重建残差覆盖”的方式,在 V=151936 下“made verification 4 to 6x slower than the stock sampler and allocated about 1 GiB transiently at 256 requests”,最终改为在内核里直接抽 key-B token。
4) 独立性能验证(tomasruizt):Qwen3.5-9B + MTP 下接受长度变化 -0.033%(p=0.82)、TPOT +1.72%(p=4.9e-8)、TTFT -2.33%,torch profile 中“No CPU-GPU syncs found”。

实现拆解

  1. 配置与算法契约(vllm/config/watermarking.py、vllm/config/vllm.py):算法枚举扩展为 gumbel | dual_key_gumbel,新增 alpha(默认 0.1)、deduplicate_contexts(none|single_turn|all,默认 single_turn)、deduplicate_contexts_max_history(默认 8192,可置 None)、allow_target_only_watermarking;新增 derive_watermark_key(key, domain) 用 SHA-256 做域分离,从同一主 key 派生 key_a/key_b;supports_speculative_decoding 从静态字典改为按算法判断;校验器由 validate_key 重命名为 validate_watermark_settings,并在关闭去重或历史长度低于 _MIN_RECOMMENDED_DEDUP_HISTORY=1024 时对两种 Gumbel 算法统一发出退化生成告警。跨配置校验补充了 dual_key 需要概率式 draft_sample_method、alpha 在 target-only 模式下未被使用等约束。

  2. 采样主干抽象(vllm/v1/watermarking/watermarker.py、gumbel.py):RandomSampler 从 Callable 类型别名改为携带 expanded_idx_mapping/temperatures/seeds/positions/use_fp64 的 dataclass,其 call 内部走 gumbel_sample;Watermarker.sample 变成模板方法——先尝试 _try_sample_mixed 在算法内部完成“水印 + 非水印”融合采样,否则调用抽象方法 _sample_watermarked,再用 skip_mask 做 torch.where 回退。原 GumbelWatermarker.sample 因此拆成 _sample_watermarked/_try_sample_mixed。新增 DualKeyGumbelWatermarker(组合两个 GumbelWatermarker,按 alpha 用随机流在 key A/B 间路由)与 DualKeyGumbelWatermarkDetector(两 key 分数加权早融合 + 按 variance 重标定 p 值)。

  3. 投机解码接入(vllm/v1/watermarking/spec_decode.py 新文件、v1/worker/gpu/spec_decode/rejection_sampler_utils.py、rejection_sampler.py、speculator.py):新模块用 create_speculative_target_watermarker/create_speculative_draft_watermarker 把双 key 容器拆成目标侧与草稿侧;DraftWatermarker 为草稿模型维护每请求的 context 窗口与 enabled 开关;_resolve_watermark_key 显式拒绝把未拆分的容器直接用于恢复采样,避免用错 key。rejection_sampler_utils 的 _resample_kernel 新增 WATERMARK/CONTEXT_WIDTH/contexts/watermarking/watermark_key 参数,在 kernel 内调用 philox_gumbel_block_argmax 直接抽 key-B 的残差与 bonus token;rejection_sampler.py 增加 _watermarking_kwargs,speculator.py 增加 prepare_watermarking 喂入上下文和开关。

  4. 生成期上下文去重(vllm/v1/worker/gpu/sample/watermark.py、vllm/v1/watermarking/gpu_sampler.py):新增 repeated_context_mask 入口与 _repeated_context_mask_kernel(Triton)/_repeated_context_mask_cpu(参考实现),扫描 all_token_ids 判断当前 context 是否已经在该请求历史中出现;gpu_sampler 的 _sample_random 在 dedup 不为 none 时计算 repeated_contexts 并折叠进 skip_mask,同时把采样状态打包成 RandomSampler 传给水印器;_get_contexts 新增 expanded_local_pos/draft_sampled 形参,用于在投机验证的扁平化行上区分已提交 token 与草稿 token,正确拼出每个待采样位置的 context。

  5. 预热与测试/CI 配套:新文件 vllm/model_executor/warmup/watermark_sample_warmup.py 在启动时按 (key, logits dtype, skip_mask, USE_FP64) 组合预热 _philox_gumbel_kernel,避免首个普通请求承担多秒编译(jit_monitor_mode=error 下还会直接使引擎退出);spec_decode_rejection_warmup.py 扩展到 watermark 参数。测试侧新增 tests/model_executor/test_watermark_sample_warmup.py、tests/model_executor/test_spec_decode_rejection_warmup.py,扩展 tests/watermarking/test_watermarking.py、tests/watermarking/test_gumbel.py、tests/v1/spec_decode/test_rejection_sampler_utils.py、tests/test_config.py;CI 增加 DSpark + TP2 下 dual-key 水印的 GSM8K 准确率与草稿接受率用例。

文件 模块 状态 重要度
vllm/v1/watermarking/spec_decode.py 水印投机解码 added 8.82
vllm/v1/watermarking/gumbel.py 水印算法 modified 8.55
vllm/v1/worker/gpu/sample/watermark.py 水印采样内核 modified 8.65
vllm/v1/watermarking/watermarker.py 水印接口 modified 8.09
vllm/v1/watermarking/gpu_sampler.py 水印采样器 modified 7.97
vllm/config/watermarking.py 水印配置 modified 7.92
vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py 拒绝采样 modified 7.54
vllm/model_executor/warmup/watermark_sample_warmup.py 水印预热 added 8.66
vllm/model_executor/warmup/spec_decode_rejection_warmup.py 拒绝采样预热 modified 6.9
vllm/v1/worker/gpu/model_runner.py 模型执行器 modified 6.46
vllm/v1/worker/gpu/spec_decode/speculator.py 投机解码器 modified 6.38
tests/model_executor/test_watermark_sample_warmup.py 预热测试 added 7.74

关键符号

DraftWatermarker.sample DraftWatermarker.prepare create_speculative_target_watermarker create_speculative_draft_watermarker watermarked_rejection_sample _resolve_watermark_key DualKeyGumbelWatermarker.sample DualKeyGumbelWatermarker._routing_logits DualKeyGumbelWatermarkDetector._get_p_value GumbelWatermarker._try_sample_mixed Watermarker.sample RandomSampler.__call__ repeated_context_mask _repeated_context_mask_kernel philox_gumbel_block_argmax GPUWatermarkSampler._get_repeated_contexts GPUWatermarkSampler._get_contexts derive_watermark_key validate_watermark_settings watermark_sample_warmup _philox_sampler_keys _seeded_resample_argmax

分析完成后,这里会展示 LLM 生成的相对完整源码片段和详细注释。

评论区精华

水印类层次是否该扁平化 设计

tomasruizt 指出“it is not super easy to understand the relationship of different watermarking classes and their specific responsibilities”,列出三条简化建议:去掉 WatermarkedRejectionSampler 子类、给 DualKeyGumbelWatermarker 改用组合、把 DraftWatermarker 的 contexts/enabled 抽成状态对象;TQCB 接受前两条,但认为 DraftWatermarker 承载的状态难以抽离,否则“without a tonne of duplication”。

结论:tomasruizt 直接提交补丁 25df1e46(删除子类,改为向 RejectionSampler 传可选 target watermark key)与 1153b1c0(DualKeyGumbelWatermarker 改用显式组合),均被合入;DraftWatermarker 与协议边界保留。 · 已解决

27B 模型检测率偏低的归因 正确性

simon-veitner-redhat 复现了 23% 的检测率后判断“the cause is how little entropy the 27B leaves in five-shot GSM8K answers, not the dual-key scheme, MTP, or the base feature”,并补充“On ordinary prose the 27B is detected in 99–100% of completions. Not a blocker; the description and docs should say this instead of guessing at it”。

结论:不改算法,改为在描述与文档中说明检测率依赖输出熵(并附上 p(x) > 0.999 的 token 占比等证据)。 · 已解决

context mask kernel 的历史扫描越界读 正确性

simon-veitner-redhat 说明 _repeated_context_mask_kernel 的下界 mask 缺失导致每个扫描的最后一块最多多读 BLOCK-3 个位置,请求位于最后一个槽位时会越过 all_token_ids,compute-sanitizer 报 47 次非法全局读:“the comparison already excluded those lanes, so results were unaffected; the read was the defect”。

结论:循环起点对齐 BLOCK 边界并用统一的 in_window 谓词裁剪,同时把 context token 改成标量加载、去掉 CONTEXT_BLOCK。 · 已解决

key-B 恢复采样的性能返工 性能

最初在 PyTorch 里重算残差覆盖 kernel 结果,d01c3916 的 commit message 记录其在 V=151936 下“made verification 4 to 6x slower than the stock sampler and allocated about 1 GiB transiently at 256 requests”;改为在 _resample_kernel 内用 philox_gumbel_block_argmax 直接抽 key-B token。

结论:恢复与 bonus token 的 key-B 采样全部进入 kernel,PyTorch 侧重建残差的路径被移除;测试用水印恢复与 philox_gumbel_sample 的等价性做回归保护。 · 已解决

水印 sampler 的 JIT 预热缺失 性能

97a26400 指出水印引擎会在首个普通请求内触发 _philox_gumbel_kernel 编译,带来多秒延迟尖峰,并在 jit_monitor_mode=error 下直接让引擎退出;通用 sampler 预热因 SamplingParams.for_sampler_warmup() 走 fp32 拷贝,覆盖不到模型 dtype 的 specialization。

结论:新增 watermark_sample_warmup,按 (key, dtype, skip_mask, USE_FP64) 预热,并在测试中逐条断言组合集合。 · 已解决

独立性能与正确性验证 性能

tomasruizt 在 H100 上跑 GSM8K 正确性与 vllm bench sweep serve:接受长度 -0.033%(p=0.82)、TPOT +1.72%(p=4.9e-8)、TTFT -2.33%(p=0.030),并提供 torch profile 说明“No CPU-GPU syncs found”。

结论:确认水印对接受率基本无影响、额外开销可控,作为合入依据。 · 已解决

CI 与缺失 watermark_config 的配置对象 测试

b9e8b1f8 记录 tests/v1/worker/test_gpu_autoregressive_speculator.py 用裸 SimpleNamespace 构造 vllm_config,新增的硬读取 watermark_config 触发 AttributeError,导致 H200/MI300 上的 CI job 失败;此前还有合并冲突与需要 cherry-pick 的 CI 修复。

结论:在三个新增读取点改用 getattr(..., "watermark_config", None),与 VllmConfig 自身的惯用法保持一致。 · 已解决

风险与影响

1) 采样热路径与语义变化:Watermarker.sample 由抽象方法变为模板方法、抽象方法改为 _sample_watermarked,任何仓库外的自定义水印器子类都会在实例化时失败;GPUWatermarkSampler.init 新增 deduplicate_contexts 等参数,属于内部 API 但在下游 fork 中常见被包装。
2) 性能:dual-key 模式每个采样步要分别用 key A、key B 各跑一次全词表 Gumbel 采样,再用 random_sampler(routing_logits) 做第三次(小词表)采样,开销随批次与词表放大;上下文去重的微基准显示小批量下 ms/step 增加 166%-208%(1 token/步时从 0.036ms 到 0.110ms),大模型端到端基本在 ±5% 以内,但这是额外成本而非免费。
3) 正确性/可检测性:去重意味着部分 token 故意不打水印,PR body 承认“a user will manipulate watermarked text to exclude an initial watermarking of a repeated context”会削弱检测信号;d7fe3812 进一步说明 deduplicate_contexts=all 的检测必须把 prompt 作为上下文前缀,否则检测口径与生成不一致(实测 74.7% vs 64.0%)。27B 在低熵输出上检测率仅 23%,用户可能误判水印失效。
4) 内核边界:_repeated_context_mask_kernel 的历史扫描窗口曾经越过 total_len(8d0f1870 修复,47 次非法全局读),MAX_HISTORY/INCLUDE_PROMPT/SKIP_PARTIAL_CONTEXT 组合增多后该类边界问题容易复发;gpu_sampler._get_contexts 的 draft_indices 计算带 clamp,仍需靠 tests/v1/spec_decode/test_rejection_sampler_utils.py 的对比测试兜底。
5) 预热覆盖:watermark_sample_warmup 依赖 sampler 的 watermarker/deduplicate_contexts/use_fp64_gumbel 三个属性,一旦新增 specialization 维度而预热未同步,会退回“首个请求承担 JIT 编译”的旧问题(97a26400 所修)。

对用户:可以在开启投机解码的同时使用水印,新增 dual_key_gumbel 算法与 alpha 旋钮;默认 deduplicate_contexts=single_turn 会改变已有水印部署的采样行为(重复 context 回退普通采样),关闭去重会收到退化生成告警。对系统:改动落在 V1 GPU runner 的采样与验证热路径、Triton kernel、以及启动期 JIT 预热,影响面覆盖 MRv2、spec decode、水印三个子系统;spec_decode.py 与 rejection kernel 的 key 路由一旦出错会静默降低可检测性或改变接受率,因此配套测试较重。对团队:这是一次跨公司协作的典型案例(Mistral 的 TQCB/raphael.rialland、Red Hat 的 simon-veitner-redhat、以及 tomasruizt 直接提交重构补丁),review 中的抽象简化与内核修复都被实际合入。

核心采样路径变更 水印器基类 API 破坏性调整 Triton 内核越界读(已修复,需防回归) 双 key 采样引入额外 kernel 启动 检测率依赖模型输出熵 配置校验行为变化

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接,后续同步到相关引用后会出现在这里。

完整报告

参与讨论