[Rust Frontend] Forward per-request watermarking controls (#56338)
Co-authored-by: Codex <noreply@openai.com> Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>
Nicolò Lucchesi committed
Sep 14, 2026 at 17:17 UTC
1713a094915d643c4ffaa25f421353108ee2049f
14 files changed
+133
-1
docs/features/watermarking.md
+6
-1
@@ -29,7 +29,12 @@ from vllm import SamplingParams
29
sampling_params = SamplingParams(watermarking=False)
30
```
31
32
-The OpenAI-compatible APIs accept the same `watermarking: false` request field.
32
+The Python OpenAI-compatible APIs and the Rust chat/completions APIs accept the
33
+same `watermarking: false` request field. The Rust token API accepts it in
34
+`sampling_params`, and the Rust gRPC `GenerateRequest` accepts
35
+the optional `watermarking` field. Omitting the field keeps watermarking enabled
36
+when the engine has a watermark configuration.
37
+
38
Deployments that require watermarking must restrict this field to trusted
39
callers, or strip and validate it at the ingress boundary, so untrusted clients
40
cannot opt out.
rust/proto/inference.proto
+3
@@ -53,6 +53,9 @@ message GenerateRequest {
53
54
// Loaded LoRA name; empty selects the base model.
55
string lora_name = 15;
56
+
57
+ // Apply the engine's configured watermark. Defaults to true when omitted.
58
+ optional bool watermarking = 16;
59
}
60
61
message RandomSampling {
rust/src/engine-core-client/src/protocol/sampling.rs
+9
@@ -20,6 +20,10 @@ fn default_temperature() -> f32 {
20
1.0
21
}
22
23
+fn default_watermarking() -> bool {
24
+ true
25
+}
26
+
27
fn default_max_tokens() -> u32 {
28
16
29
}
@@ -68,6 +72,9 @@ pub struct EngineCoreSamplingParams {
72
/// greedy sampling.
73
#[serde(default = "default_temperature")]
74
pub temperature: f32,
75
+ /// Whether to apply the engine's configured watermark to this request.
76
+ #[serde(default = "default_watermarking")]
77
+ pub watermarking: bool,
78
/// Cumulative probability threshold for nucleus sampling.
79
#[serde(default = "default_top_p")]
80
pub top_p: f32,
@@ -155,6 +162,7 @@ impl EngineCoreSamplingParams {
162
pub fn for_test() -> Self {
163
Self {
164
temperature: 1.0,
165
+ watermarking: true,
166
top_p: 1.0,
167
top_k: 0,
168
seed: None,
@@ -229,6 +237,7 @@ mod tests {
237
238
// Omitted fields -> Python defaults.
239
assert_eq!(sampling.temperature, 1.0);
240
+ assert!(sampling.watermarking);
241
assert_eq!(sampling.top_p, 1.0);
242
assert_eq!(sampling.top_k, 0);
243
assert_eq!(sampling.seed, None);
rust/src/engine-core-client/src/tests/client.rs
+2
@@ -151,6 +151,7 @@ fn sample_request_with_id(request_id: &str) -> EngineCoreRequest {
151
prompt_token_ids: Some(vec![11, 22]),
152
sampling_params: Some(EngineCoreSamplingParams {
153
temperature: 0.8,
154
+ watermarking: false,
155
top_p: 0.9,
156
top_k: 8,
157
max_tokens: 32,
@@ -2662,6 +2663,7 @@ fn python_msgpack_fixtures_match_rust_encoding() {
2663
sampling,
2664
EngineCoreSamplingParams {
2665
temperature: 1.0,
2666
+ watermarking: true,
2667
top_p: 1.0,
2668
top_k: 0,
2669
seed: None,
rust/src/engine-core-client/src/tests/python_compat.py
+2
@@ -34,6 +34,7 @@ class FinishReason(IntEnum):
34
# Mirror of real SamplingParams; omit_defaults makes fixtures match real maps.
35
class EngineCoreSamplingParams(msgspec.Struct, dict=True, omit_defaults=True):
36
temperature: float = 1.0
37
+ watermarking: bool = True
38
top_p: float = 1.0
39
top_k: int = 0
40
seed: int | None = None
@@ -123,6 +124,7 @@ request = EngineCoreRequest(
124
mm_features=None,
125
sampling_params=EngineCoreSamplingParams(
126
temperature=0.8,
127
+ watermarking=False,
128
top_p=0.9,
129
top_k=8,
130
seed=None,
rust/src/server/src/grpc/convert.rs
+22
@@ -65,6 +65,7 @@ pub fn to_text_request(
65
66
let mut sampling_params =
67
build_sampling_params(req.temperature, sampling, decoding, stopping, response)?;
68
+ sampling_params.watermarking = req.watermarking.unwrap_or(true);
69
70
// Thread KVCacheParameters → SamplingParams fields.
71
if let Some(kv) = kv {
@@ -521,6 +522,7 @@ impl ResponseOpts {
522
523
#[cfg(test)]
524
mod tests {
525
+ use prost::Message as _;
526
use vllm_engine_core_client::protocol::output::StopReason;
527
use vllm_text::{
528
FinishReason, Finished, Prompt, SamplingHints, SamplingLimits, lower_sampling_params,
@@ -539,6 +541,26 @@ mod tests {
541
}
542
}
543
544
+ #[test]
545
+ fn watermarking_defaults_and_opt_out_survive_protobuf_conversion() {
546
+ for watermarking in [None, Some(true), Some(false)] {
547
+ let request = pb::GenerateRequest {
548
+ watermarking,
549
+ ..base_request()
550
+ };
551
+ let encoded = request.encode_to_vec();
552
+ for stream in [false, true] {
553
+ let decoded = pb::GenerateRequest::decode(encoded.as_slice()).unwrap();
554
+ let text = to_text_request(decoded, stream, &["test-model".to_string()])
555
+ .expect("convert request");
556
+ assert_eq!(
557
+ text.sampling_params.watermarking,
558
+ watermarking.unwrap_or(true)
559
+ );
560
+ }
561
+ }
562
+ }
563
+
564
#[test]
565
fn temperature_propagates_from_top_level_request_field() {
566
let req = pb::GenerateRequest {
rust/src/server/src/routes/inference/generate/convert.rs
+24
@@ -164,6 +164,30 @@ mod tests {
164
);
165
}
166
167
+ #[test]
168
+ fn prepare_generate_request_preserves_watermarking_defaults_and_opt_out() {
169
+ for watermarking in [None, Some(true), Some(false)] {
170
+ let mut body = json!({
171
+ "token_ids": [11, 22],
172
+ "sampling_params": {}
173
+ });
174
+ if let Some(watermarking) = watermarking {
175
+ body["sampling_params"]["watermarking"] = json!(watermarking);
176
+ }
177
+ let prepared = prepare_generate_request(
178
+ serde_json::from_value(body).expect("parse request"),
179
+ &served(&["test-model"]),
180
+ ResolvedRequestContext::default(),
181
+ None,
182
+ )
183
+ .expect("prepare request");
184
+ assert_eq!(
185
+ prepared.text_request.sampling_params.watermarking,
186
+ watermarking.unwrap_or(true)
187
+ );
188
+ }
189
+ }
190
+
191
#[test]
192
fn prepare_generate_request_forwards_thinking_token_budget() {
193
let request: GenerateRequest = serde_json::from_value(json!({
rust/src/server/src/routes/openai/chat_completions/convert.rs
+1
@@ -150,6 +150,7 @@ pub(super) fn prepare_chat_request(
150
messages,
151
sampling_params: SamplingParams {
152
temperature: request.temperature,
153
+ watermarking: request.watermarking,
154
top_p: request.top_p,
155
top_k: request.top_k,
156
seed: request.seed,
rust/src/server/src/routes/openai/chat_completions/types.rs
+5
@@ -116,6 +116,10 @@ pub struct ChatCompletionRequest {
116
pub user: Option<String>,
117
118
// -------- vLLM Sampling Parameters --------
119
+ /// Whether to apply the engine's configured watermark to this request.
120
+ #[serde(default = "default_true")]
121
+ pub watermarking: bool,
122
+
123
/// Use beam search instead of sampling
124
#[serde(default)]
125
pub use_beam_search: bool,
@@ -282,6 +286,7 @@ impl Default for ChatCompletionRequest {
286
include_reasoning: true,
287
parallel_tool_calls: None,
288
user: None,
289
+ watermarking: true,
290
use_beam_search: false,
291
top_k: None,
292
min_p: None,
rust/src/server/src/routes/openai/completions/convert.rs
+1
@@ -115,6 +115,7 @@ pub(super) fn prepare_completion_request(
115
mm_features: None,
116
sampling_params: SamplingParams {
117
temperature: request.temperature,
118
+ watermarking: request.watermarking,
119
top_p: request.top_p,
120
top_k: request.top_k,
121
seed: request.seed,
rust/src/server/src/routes/openai/completions/types.rs
+4
@@ -90,6 +90,10 @@ pub struct CompletionRequest {
90
pub user: Option<String>,
91
92
// -------- vLLM Sampling Parameters --------
93
+ /// Whether to apply the engine's configured watermark to this request.
94
+ #[serde(default = "default_true")]
95
+ pub watermarking: bool,
96
+
97
/// Options for streaming response
98
pub stream_options: Option<StreamOptions>,
99
rust/src/server/src/routes/tests.rs
+43
@@ -2736,6 +2736,49 @@ async fn non_stream_chat_returns_json_response() {
2736
assert!(!message.contains_key("tool_calls"), "{json}");
2737
}
2738
2739
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2740
+#[serial]
2741
+async fn openai_watermarking_defaults_and_opt_out_reach_engine_core() {
2742
+ for endpoint in ["/v1/chat/completions", "/v1/completions"] {
2743
+ for stream in [false, true] {
2744
+ for watermarking in [None, Some(true), Some(false)] {
2745
+ let (mut app, engine_task) = test_app_with_backend_and_engine_request_check(
2746
+ Arc::new(FakeChatBackend::new()),
2747
+ move |request| {
2748
+ let params = request.sampling_params.as_ref().expect("sampling params");
2749
+ assert_eq!(params.watermarking, watermarking.unwrap_or(true));
2750
+ },
2751
+ )
2752
+ .await;
2753
+
2754
+ let mut body = json!({"stream": stream, "max_tokens": 2});
2755
+ if endpoint == "/v1/chat/completions" {
2756
+ body["messages"] = json!([{"role": "user", "content": "hi"}]);
2757
+ } else {
2758
+ body["prompt"] = json!("hi");
2759
+ }
2760
+ if let Some(watermarking) = watermarking {
2761
+ body["watermarking"] = json!(watermarking);
2762
+ }
2763
+ let response = app
2764
+ .call(
2765
+ Request::builder()
2766
+ .method("POST")
2767
+ .uri(endpoint)
2768
+ .header("content-type", "application/json")
2769
+ .body(Body::from(body.to_string()))
2770
+ .expect("build request"),
2771
+ )
2772
+ .await
2773
+ .expect("call app");
2774
+ assert_eq!(response.status(), StatusCode::OK);
2775
+ to_bytes(response.into_body(), usize::MAX).await.expect("drain response");
2776
+ engine_task.await.expect("mock engine task");
2777
+ }
2778
+ }
2779
+ }
2780
+}
2781
+
2782
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2783
#[serial]
2784
async fn non_stream_chat_image_url_reaches_engine_mm_features() {
rust/src/text/src/lower.rs
+8
@@ -92,6 +92,7 @@ pub fn lower_sampling_params(
92
) -> Result<EngineCoreSamplingParams> {
93
let SamplingParams {
94
temperature,
95
+ watermarking,
96
top_p,
97
top_k,
98
seed,
@@ -164,6 +165,7 @@ pub fn lower_sampling_params(
165
166
let params = EngineCoreSamplingParams {
167
temperature,
168
+ watermarking,
169
top_p,
170
top_k,
171
seed,
@@ -632,6 +634,7 @@ mod tests {
634
expect_test::expect![[r#"
635
EngineCoreSamplingParams {
636
temperature: 1.0,
637
+ watermarking: true,
638
top_p: 1.0,
639
top_k: 0,
640
seed: None,
@@ -686,6 +689,7 @@ mod tests {
689
expect_test::expect![[r#"
690
EngineCoreSamplingParams {
691
temperature: 1.0,
692
+ watermarking: true,
693
top_p: 1.0,
694
top_k: 0,
695
seed: None,
@@ -854,6 +858,7 @@ mod tests {
858
expect_test::expect![[r#"
859
EngineCoreSamplingParams {
860
temperature: 0.6,
861
+ watermarking: true,
862
top_p: 0.95,
863
top_k: 20,
864
seed: None,
@@ -918,6 +923,7 @@ mod tests {
923
expect_test::expect![[r#"
924
EngineCoreSamplingParams {
925
temperature: 1.0,
926
+ watermarking: true,
927
top_p: 1.0,
928
top_k: 0,
929
seed: None,
@@ -990,6 +996,7 @@ mod tests {
996
expect_test::expect![[r#"
997
EngineCoreSamplingParams {
998
temperature: 0.2,
999
+ watermarking: true,
1000
top_p: 0.3,
1001
top_k: 4,
1002
seed: None,
@@ -1240,6 +1247,7 @@ mod tests {
1247
expect_test::expect![[r#"
1248
EngineCoreSamplingParams {
1249
temperature: 0.8,
1250
+ watermarking: true,
1251
top_p: 0.9,
1252
top_k: 12,
1253
seed: None,
rust/src/text/src/request.rs
+3
@@ -69,6 +69,8 @@ pub struct SamplingParams {
69
/// Controls randomness. Lower values are more deterministic; zero means
70
/// greedy sampling. `None` means no explicit user override.
71
pub temperature: Option<f32>,
72
+ /// Whether to apply the engine's configured watermark to this request.
73
+ pub watermarking: bool,
74
/// Cumulative probability threshold for nucleus sampling.
75
pub top_p: Option<f32>,
76
/// Maximum number of top tokens to consider. `Some(0)` means all tokens.
@@ -144,6 +146,7 @@ impl Default for SamplingParams {
146
fn default() -> Self {
147
Self {
148
temperature: None,
149
+ watermarking: true,
150
top_p: None,
151
top_k: None,
152
seed: None,