From 14e3902d5136e2fb00ef86c05fdc8d098f736283 Mon Sep 17 00:00:00 2001
From: Alec <35311602+alec-flowers@users.noreply.github.com>
Date: Wed, 7 Oct 2026 07:46:44 -0700
Subject: [PATCH] [Frontend][Rust] Add gRPC forbidden token sequences and cache
 usage (#59837)

Signed-off-by: Alec Flowers <aflowers@nvidia.com>
---
 rust/proto/inference.proto                    |  4 +
 rust/src/server/src/grpc/convert.rs           | 46 +++++++++++-
 rust/src/server/src/grpc/tests.rs             |  2 +
 .../routes/openai/chat_completions/convert.rs |  1 +
 .../src/routes/openai/completions/convert.rs  |  1 +
 rust/src/text/src/lower.rs                    | 75 ++++++++++++++++---
 rust/src/text/src/lower/token_ids.rs          |  5 ++
 rust/src/text/src/request.rs                  |  3 +
 8 files changed, 125 insertions(+), 12 deletions(-)

diff --git a/rust/proto/inference.proto b/rust/proto/inference.proto
index 2bc35bcb64..8189758212 100644
--- a/rust/proto/inference.proto
+++ b/rust/proto/inference.proto
@@ -102,6 +102,9 @@ message DecodingParameters {
     bool json_object = 10;
     string structural_tag = 11;
   }
+
+  // Token sequences that must not be generated. Each sequence must be nonempty.
+  repeated TokenIds bad_words_token_ids = 12;
 }
 
 message StoppingCriteria {
@@ -210,6 +213,7 @@ message FinishInfo {
   google.protobuf.Struct kv_transfer_params = 6;
   //uint64 seed = 7;
   google.protobuf.Struct ec_transfer_params = 8;
+  optional uint32 num_cached_tokens = 9;
 }
 
 // Info for candidate tokens other than the input/sampled
diff --git a/rust/src/server/src/grpc/convert.rs b/rust/src/server/src/grpc/convert.rs
index d5dd1ae64d..8f2c589f15 100644
--- a/rust/src/server/src/grpc/convert.rs
+++ b/rust/src/server/src/grpc/convert.rs
@@ -203,6 +203,10 @@ fn build_sampling_params(
         if !d.allowed_token_ids.is_empty() {
             params.allowed_token_ids = Some(d.allowed_token_ids.clone());
         }
+        if !d.bad_words_token_ids.is_empty() {
+            params.bad_words_token_ids =
+                Some(d.bad_words_token_ids.iter().map(|sequence| sequence.ids.clone()).collect());
+        }
         params.structured_outputs = convert_structured_output(d)?;
     }
 
@@ -403,6 +407,7 @@ fn to_finish_info(finished: &Finished, token_ids: &[u32]) -> Result<pb::FinishIn
         stop_reason,
         kv_transfer_params: finished.kv_transfer_params.as_ref().and_then(json_to_proto_struct),
         ec_transfer_params: finished.ec_transfer_params.as_ref().and_then(json_to_proto_struct),
+        num_cached_tokens: Some(finished.usage.cached_token_count as u32),
     })
 }
 
@@ -580,6 +585,41 @@ mod tests {
         }
     }
 
+    #[test]
+    fn grpc_bad_words_token_ids_survive_protobuf_conversion_and_lowering() {
+        let request = pb::GenerateRequest {
+            decoding: Some(pb::DecodingParameters {
+                bad_words_token_ids: vec![
+                    pb::TokenIds { ids: vec![5] },
+                    pb::TokenIds { ids: vec![7, 11] },
+                ],
+                ..Default::default()
+            }),
+            ..base_request()
+        };
+        let encoded = request.encode_to_vec();
+        for stream in [false, true] {
+            let decoded = pb::GenerateRequest::decode(encoded.as_slice()).unwrap();
+            let text = to_text_request(decoded, stream, &["test-model".to_string()]).unwrap();
+            let engine = vllm_text::lower_text_request(
+                text,
+                vec![1],
+                SamplingHints::default(),
+                SamplingLimits {
+                    max_model_len: 256,
+                    max_logprobs: 20,
+                    model_vocab_size: 512,
+                    tokenizer_vocab_size: 512,
+                },
+                &TestTokenizer::new(),
+            )
+            .unwrap()
+            .generate_request;
+            let params = engine.sampling_params;
+            assert_eq!(params.bad_words_token_ids, Some(vec![vec![5], vec![7, 11]]));
+        }
+    }
+
     #[test]
     fn watermarking_defaults_and_opt_out_survive_protobuf_conversion() {
         for watermarking in [None, Some(true), Some(false)] {
@@ -859,7 +899,9 @@ mod tests {
 
     #[test]
     fn to_sequence_output_threads_token_ids_into_eos_id() {
-        let fin = finished(FinishReason::Stop(None));
+        let mut fin = finished(FinishReason::Stop(None));
+        fin.usage.prompt_token_count = 8;
+        fin.usage.cached_token_count = 5;
         let opts = ResponseOpts {
             output_text: true,
             output_token_ids: true,
@@ -868,10 +910,12 @@ mod tests {
 
         let out = to_sequence_output("hello", &[10, 20, 30], None, Some(&fin), &opts)
             .expect("sequence output");
+        let out = pb::SequenceOutput::decode(out.encode_to_vec().as_slice()).unwrap();
 
         let finish = out.finish_info.expect("finish_info should be present");
         assert_eq!(finish.finish_reason, PbFinishReason::Stop as i32);
         assert_eq!(finish.stop_reason, Some(PbStopReason::EosTokenId(30)));
+        assert_eq!(finish.num_cached_tokens, Some(5));
     }
 
     #[test]
diff --git a/rust/src/server/src/grpc/tests.rs b/rust/src/server/src/grpc/tests.rs
index 9f01886062..d8d0fb4e6c 100644
--- a/rust/src/server/src/grpc/tests.rs
+++ b/rust/src/server/src/grpc/tests.rs
@@ -690,6 +690,7 @@ async fn unary_generate_returns_collected_text() {
         pb::finish_info::FinishReason::Stop as i32
     );
     assert_eq!(finish.num_output_tokens, 3);
+    assert_eq!(finish.num_cached_tokens, Some(0));
 
     let prompt = response.prompt_info.expect("prompt_info present");
     assert_eq!(prompt.num_prompt_tokens, 5); // "hello" = 5 bytes
@@ -1438,6 +1439,7 @@ async fn streaming_generate_yields_incremental_responses() {
         .find_map(|r| r.outputs.as_ref())
         .expect("at least one output");
     let finish = last_output.finish_info.as_ref().expect("finish_info on last output");
+    assert_eq!(finish.num_cached_tokens, Some(0));
     assert_eq!(
         finish.finish_reason,
         pb::finish_info::FinishReason::Stop as i32
diff --git a/rust/src/server/src/routes/openai/chat_completions/convert.rs b/rust/src/server/src/routes/openai/chat_completions/convert.rs
index 998ca4c939..4f80dda6c7 100644
--- a/rust/src/server/src/routes/openai/chat_completions/convert.rs
+++ b/rust/src/server/src/routes/openai/chat_completions/convert.rs
@@ -171,6 +171,7 @@ pub(super) fn prepare_chat_request(
             logit_bias: convert_logit_bias(request.logit_bias)?,
             allowed_token_ids: request.allowed_token_ids,
             bad_words: request.bad_words,
+            bad_words_token_ids: None,
             logprob_token_ids: None,
             structured_outputs,
             skip_reading_prefix_cache: None,
diff --git a/rust/src/server/src/routes/openai/completions/convert.rs b/rust/src/server/src/routes/openai/completions/convert.rs
index d05cfffc89..0b37f84057 100644
--- a/rust/src/server/src/routes/openai/completions/convert.rs
+++ b/rust/src/server/src/routes/openai/completions/convert.rs
@@ -136,6 +136,7 @@ pub(super) fn prepare_completion_request(
             logit_bias: convert_logit_bias(request.logit_bias)?,
             allowed_token_ids: request.allowed_token_ids,
             bad_words: None,
+            bad_words_token_ids: None,
             logprob_token_ids: None,
             structured_outputs,
             skip_reading_prefix_cache: None,
diff --git a/rust/src/text/src/lower.rs b/rust/src/text/src/lower.rs
index cf53704374..2b09070049 100644
--- a/rust/src/text/src/lower.rs
+++ b/rust/src/text/src/lower.rs
@@ -119,6 +119,7 @@ pub fn lower_sampling_params(
         logit_bias,
         allowed_token_ids,
         bad_words,
+        bad_words_token_ids,
         logprob_token_ids,
         structured_outputs,
         skip_reading_prefix_cache,
@@ -173,6 +174,11 @@ pub fn lower_sampling_params(
         merge_unique_token_ids(&mut stop_token_ids, extra_eos_token_ids.iter().copied());
     }
 
+    let mut bad_words_token_ids = bad_words_token_ids.unwrap_or_default();
+    if let Some(tokenized) = tokenize_bad_words(bad_words.as_deref(), tokenizer)? {
+        bad_words_token_ids.extend(tokenized);
+    }
+
     let params = EngineCoreSamplingParams {
         temperature,
         watermarking,
@@ -201,7 +207,7 @@ pub fn lower_sampling_params(
         all_stop_token_ids,
         logit_bias,
         allowed_token_ids,
-        bad_words_token_ids: tokenize_bad_words(bad_words.as_deref(), tokenizer)?,
+        bad_words_token_ids: (!bad_words_token_ids.is_empty()).then_some(bad_words_token_ids),
         // TODO: Validate structured-output schemas and regexes before submitting requests to engine-core.
         structured_outputs,
         logprob_token_ids,
@@ -1255,25 +1261,72 @@ mod tests {
     #[test]
     fn lower_sampling_params_rejects_out_of_vocab_bad_words() {
         let tokenizer = TestTokenizer::new().with_regular_token("blocked", 2000);
-        let error = lower_sampling_params(
+        for sampling_params in [
             SamplingParams {
                 bad_words: Some(vec!["blocked".to_string()]),
                 ..Default::default()
             },
-            SamplingHints::default(),
+            SamplingParams {
+                bad_words_token_ids: Some(vec![vec![1999, 2000]]),
+                ..Default::default()
+            },
+        ] {
+            let error = lower_sampling_params(
+                sampling_params,
+                SamplingHints::default(),
+                sample_sampling_limits(),
+                3,
+                &tokenizer,
+            )
+            .unwrap_err();
+
+            assert!(matches!(
+                error,
+                Error::TokenIds(TokenIdsError::OutOfVocab {
+                    parameter: "bad_words",
+                    token_ids,
+                    vocab_size: 2000,
+                }) if token_ids == vec![2000]
+            ));
+        }
+    }
+
+    #[test]
+    fn lower_sampling_params_preserves_raw_bad_words_alongside_strings() {
+        let tokenizer = TestTokenizer::new()
+            .with_regular_token("blocked", 300)
+            .with_regular_token(" blocked", 301);
+        for raw in [None, Some(vec![]), Some(vec![vec![7, 11]])] {
+            let params = lower_sampling_params(
+                SamplingParams {
+                    bad_words: Some(vec!["blocked".to_string()]),
+                    bad_words_token_ids: raw.clone(),
+                    ..Default::default()
+                },
+                SamplingHints::default(),
+                sample_sampling_limits(),
+                3,
+                &tokenizer,
+            )
+            .unwrap();
+            let expected = if raw.as_ref().is_some_and(|ids| !ids.is_empty()) {
+                vec![vec![7, 11], vec![300], vec![301]]
+            } else {
+                vec![vec![300], vec![301]]
+            };
+            assert_eq!(params.bad_words_token_ids, Some(expected));
+        }
+        let error = lower_sampling_params_with_limits(
+            SamplingParams {
+                bad_words_token_ids: Some(vec![vec![]]),
+                ..Default::default()
+            },
             sample_sampling_limits(),
-            3,
-            &tokenizer,
         )
         .unwrap_err();
-
         assert!(matches!(
             error,
-            Error::TokenIds(TokenIdsError::OutOfVocab {
-                parameter: "bad_words",
-                token_ids,
-                vocab_size: 2000,
-            }) if token_ids == vec![2000]
+            Error::TokenIds(TokenIdsError::EmptyBadWordSequence)
         ));
     }
 
diff --git a/rust/src/text/src/lower/token_ids.rs b/rust/src/text/src/lower/token_ids.rs
index 8512a830e6..51329f9290 100644
--- a/rust/src/text/src/lower/token_ids.rs
+++ b/rust/src/text/src/lower/token_ids.rs
@@ -12,6 +12,8 @@ use crate::SamplingLimits;
 pub enum TokenIdsError {
     #[error("allowed_token_ids should not be empty")]
     EmptyAllowedTokenIds,
+    #[error("bad_words_token_ids must not contain empty sequences")]
+    EmptyBadWordSequence,
     #[error(
         "token_id(s) {token_ids:?} in {parameter} are out of vocabulary. \
          Vocabulary size: {vocab_size}"
@@ -96,6 +98,9 @@ pub(crate) fn validate_vocab_range(
     }
 
     if let Some(bad_words_token_ids) = params.bad_words_token_ids.as_deref() {
+        if bad_words_token_ids.iter().any(Vec::is_empty) {
+            return Err(TokenIdsError::EmptyBadWordSequence);
+        }
         validate_param(
             "bad_words",
             bad_words_token_ids.iter().flatten().copied(),
diff --git a/rust/src/text/src/request.rs b/rust/src/text/src/request.rs
index c37e21060e..08758748ac 100644
--- a/rust/src/text/src/request.rs
+++ b/rust/src/text/src/request.rs
@@ -129,6 +129,8 @@ pub struct SamplingParams {
     pub allowed_token_ids: Option<Vec<u32>>,
     /// Words to avoid during generation (tokenized to IDs during lowering).
     pub bad_words: Option<Vec<String>>,
+    /// Pre-tokenized sequences to prohibit, in addition to `bad_words`.
+    pub bad_words_token_ids: Option<Vec<Vec<u32>>>,
     /// Specific token IDs for which log probabilities should be returned at
     /// each position.
     ///
@@ -173,6 +175,7 @@ impl Default for SamplingParams {
             logit_bias: None,
             allowed_token_ids: None,
             bad_words: None,
+            bad_words_token_ids: None,
             logprob_token_ids: None,
             structured_outputs: None,
             skip_reading_prefix_cache: None,