[Frontend][Rust] Add gRPC forbidden token sequences and cache usage (#59837)

Signed-off-by: Alec Flowers <aflowers@nvidia.com>

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)
         ));
     }
 

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(),

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,