iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

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

Browse Source 
Signed-off-by: Alec Flowers <aflowers@nvidia.com>
This commit is contained in:
Alec
2026-10-07 07:46:44 -07:00
committed by GitHub
parent 548597367e
commit 14e3902d51
8 changed files with 125 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rust/proto/inference.proto
+4
View File
@@ -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
rust/src/server/src/grpc/convert.rs
+45 -1
View File
@@ -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]
rust/src/server/src/grpc/tests.rs
+2
View File
@@ -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
rust/src/server/src/routes/openai/chat_completions/convert.rs
+1
View File
@@ -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,
rust/src/server/src/routes/openai/completions/convert.rs
+1
View File
@@ -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,
rust/src/text/src/lower.rs
+64 -11
View File
@@ -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
+5
View File
@@ -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
+3
View File
@@ -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,