Function apply_repeat_penalty_gpu

cake-core/src/models/common/text_model.rs:60–99

Apply repeat penalty entirely on GPU to avoid costly GPU↔CPU round-trips. The upstream `candle_transformers::utils::apply_repeat_penalty` copies the entire logits tensor (vocab_size × 4 bytes ≈ 600 KB) to CPU, modifies a handful of elements, then copies everything back. This forces a full GPU synchronisation and two large PCIe transfers per token. This implementation stays on-device: it selects

(
    logits: &Tensor,
    penalty: f32,
    context: &[u32],
)

/// This implementation stays on-device: it selects only the penalty positions, computes
/// sign-aware multipliers, and scatters the deltas back with `index_add`.
fn apply_repeat_penalty_gpu(
    logits: &Tensor,
    penalty: f32,
    context: &[u32],
) -> Result<Tensor> {
    // Deduplicate tokens (same semantics as the upstream version).
    let mut seen = HashSet::new();
    let unique: Vec<u32> = context
        .iter()
        .filter(|t| seen.insert(**t))
        .copied()
        .collect();

    if unique.is_empty() {
        return Ok(logits.clone());
    }

    let device = logits.device();
    let dtype = logits.dtype();
    let indices = Tensor::new(unique.as_slice(), device)?;

    // Gather logits at penalty positions  (N elements, tiny).
    let selected = logits.index_select(&indices, 0)?;

    // Sign-aware multiplier:  1/penalty for logits ≥ 0,  penalty for logits < 0.
    let is_non_negative = selected.ge(0f32)?;
    let recip = Tensor::new(1.0f32 / penalty, device)?
        .to_dtype(dtype)?
        .broadcast_as(selected.shape())?;
    let pen = Tensor::new(penalty, device)?
        .to_dtype(dtype)?
        .broadcast_as(selected.shape())?;
    let mult = is_non_negative.where_cond(&recip, &pen)?;

    // delta = selected * mult - selected  (what to add to the original logits).
    let penalized = (&selected * &mult)?;
    let delta = (&penalized - &selected)?;

    Ok(logits.index_add(&indices, &delta, 0)?)
}

/// Create the logit sampling logic from the context.
pub fn create_logits_processor(ctx: &Context) -> LogitsProcessor {