Return f32 from apply_frequency_presence_penalty, don't round back down

Codex review: the function computed the penalty in f32 but cast the
result back to the caller's original dtype before returning. For
f16/bf16 logits this could round a small penalty away outright (or
noticeably distort it) before LogitsProcessor::sample converts back to
f32 anyway for softmax. utils::apply_repeat_penalty already avoids this
by always returning f32; do the same here for consistency, including
in the zero-penalty no-op path so the function's return dtype is
uniform regardless of whether a penalty was actually applied.

Adds a regression test with a f16 logit large enough (1000.0, ULP
~0.98) that a 0.5 penalty would be rounded away by a cast back to f16.

Author: claude

candle-transformers/src/generation/mod.rs

@@ -26,10 +26,9 @@ pub fn apply_frequency_presence_penalty(
     presence_penalty: f32,
 ) -> Result<Tensor> {
     if frequency_penalty == 0. && presence_penalty == 0. {
-        return Ok(logits.clone());
+        return logits.to_dtype(DType::F32);
     }
     let device = logits.device();
-    let dtype = logits.dtype();
     let mut logits_v = logits.to_dtype(DType::F32)?.to_vec1::<f32>()?;
     let mut counts = std::collections::HashMap::new();
     for token_id in context {
@@ -41,7 +40,11 @@ pub fn apply_frequency_presence_penalty(
         }
     }
     let logits_len = logits_v.len();
-    Tensor::from_vec(logits_v, logits_len, device)?.to_dtype(dtype)
+    // Stay in f32: rounding the adjusted logits back down to the caller's original dtype
+    // (e.g. f16/bf16) can erase a small penalty outright before `LogitsProcessor::sample`
+    // converts back to f32 anyway. Mirrors `utils::apply_repeat_penalty`, which never casts
+    // back either.
+    Tensor::from_vec(logits_v, logits_len, device)
 }
 
 pub struct LogitsProcessor {

candle-transformers/tests/generation_tests.rs

@@ -1,4 +1,4 @@
-use candle::{Device, Result, Tensor};
+use candle::{DType, Device, Result, Tensor};
 use candle_transformers::generation::{
     apply_frequency_presence_penalty, LogitsProcessor, Sampling,
 };
@@ -147,6 +147,18 @@ fn frequency_presence_penalty_noop_and_out_of_range() -> Result<()> {
     Ok(())
 }
 
+#[test]
+fn frequency_presence_penalty_preserves_precision_for_f16_input() -> Result<()> {
+    // A f16 logit of 1000.0 has an ULP of ~0.98, so subtracting a small penalty like 0.5 and
+    // rounding back down to f16 would erase it outright. Staying in f32 (matching
+    // `utils::apply_repeat_penalty`'s convention) preserves the adjustment exactly.
+    let logits = Tensor::new(&[1000.0f32], &Device::Cpu)?.to_dtype(DType::F16)?;
+    let penalized = apply_frequency_presence_penalty(&logits, &[0], 0.5, 0.0)?;
+    assert_eq!(penalized.dtype(), DType::F32);
+    assert_eq!(penalized.to_vec1::<f32>()?, [999.5]);
+    Ok(())
+}
+
 #[test]
 fn sample_gumbel() -> Result<()> {
     let mut logits_process = LogitsProcessor::from_sampling(