Skip to content

Commit 61f5623

Browse files
authored
[Sampler] Skip top-p renormalization if top-p is 1 in CPUSampler (#2528)
This PR adds a shortcut in the top-p renormalization in CPU sampler, which skips the renormalization when top-p is 1.0.
1 parent 1881992 commit 61f5623

File tree

1 file changed

+5
-0
lines changed

1 file changed

+5
-0
lines changed

cpp/serve/sampler/cpu_sampler.cc

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -178,6 +178,11 @@ void RenormalizeProbByTopP(NDArray prob, int unit_offset, double top_p, double e
178178
ICHECK(prob.DataType() == DataType::Float(32));
179179
ICHECK_EQ(prob->device.device_type, DLDeviceType::kDLCPU);
180180

181+
if (top_p == 1.0) {
182+
// No renormalization is needed if top_p is 1.
183+
return;
184+
}
185+
181186
int vocab_size = prob->shape[prob->ndim - 1];
182187
float* __restrict p_prob =
183188
static_cast<float*>(__builtin_assume_aligned(prob->data, 4)) + (unit_offset * vocab_size);

0 commit comments

Comments
 (0)