Skip to content

Commit 41819d7

Browse files
committed
fix(trtllm): fix do_sample being ignored
Currently, the do_sample option is ignored and the executor will always sample. Set top_k to 1 if do_sample is false.
1 parent 56dd0a0 commit 41819d7

File tree

1 file changed

+6
-1
lines changed

1 file changed

+6
-1
lines changed

backends/trtllm/src/looper.rs

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -98,12 +98,17 @@ fn executor_status_looper(
9898
let generation_params = &request.parameters;
9999
let stopping_params = &request.stopping_parameters;
100100
let input_ids = request.input_ids.as_deref();
101+
let top_k = if generation_params.do_sample {
102+
generation_params.top_k
103+
} else {
104+
1
105+
};
101106

102107
// Submit to the TensorRT-LLM executor for scheduling
103108
match backend.pin_mut().submit(
104109
&input_ids.unwrap(), // This is checked beforehand in validate()
105110
stopping_params.max_new_tokens,
106-
generation_params.top_k,
111+
top_k,
107112
generation_params.top_p,
108113
generation_params.temperature,
109114
generation_params.repetition_penalty,

0 commit comments

Comments
 (0)