Add Qwen3.6 MoE inference support

crates/xserv-model/src/sampling.rs

@@ -1,5 +1,6 @@
 use half::bf16;
 use rand::Rng;
+use std::collections::HashSet;
 use xserv_tensor::{DType, Device, Tensor};
 
 #[derive(Clone)]
@@ -7,6 +8,8 @@ pub struct SamplingParams {
     pub temperature: f32,
     pub top_k: usize,
     pub top_p: f32,
+    pub presence_penalty: f32,
+    pub repetition_penalty: f32,
 }
 
 impl Default for SamplingParams {
@@ -15,6 +18,8 @@ impl Default for SamplingParams {
             temperature: 0.0,
             top_k: 0,
             top_p: 1.0,
+            presence_penalty: 0.0,
+            repetition_penalty: 1.0,
         }
     }
 }
@@ -22,12 +27,21 @@ impl Default for SamplingParams {
 /// Sample a token from logits with shape [seq_len, vocab_size].
 /// Uses the last position's logits. Handles both F32 and BF16 dtypes.
 pub fn sample(logits: &Tensor, params: &SamplingParams) -> u32 {
+    sample_with_history(logits, params, &[])
+}
+
+/// Sample while applying penalties to tokens already present in the request.
+/// Presence penalty follows the OpenAI convention (subtract once per distinct
+/// token); repetition penalty follows the HF convention.
+pub fn sample_with_history(logits: &Tensor, params: &SamplingParams, history: &[u32]) -> u32 {
     assert_eq!(logits.ndim(), 2);
     // Greedy fast path: GPU argmax + 4-byte D2H instead of copying the whole
     // [seq, vocab] logits to the host and scanning it (~201k bf16/token).
     // NaN logits lose every `>` comparison in the kernel, matching the
     // NaN-safe host argmax below.
     if params.temperature == 0.0
+        && params.presence_penalty == 0.0
+        && params.repetition_penalty == 1.0
         && logits.dtype() == DType::BF16
         && matches!(logits.device(), Device::Cuda(_))
         && logits.is_contiguous()
@@ -55,11 +69,6 @@ pub fn sample(logits: &Tensor, params: &SamplingParams) -> u32 {
         _ => panic!("unsupported dtype for sampling: {:?}", logits.dtype()),
     };
 
-    // Greedy
-    if params.temperature == 0.0 {
-        return argmax(&last_row);
-    }
-
     // NaN-safe: sampling path uses partial_cmp().unwrap() in top-k/top-p
     // sorts and softmax; a single NaN logit would panic the engine thread.
     // Replace NaN with -inf (equivalent to masking) instead.
@@ -74,6 +83,29 @@ pub fn sample(logits: &Tensor, params: &SamplingParams) -> u32 {
         eprintln!("[sampling] WARNING: NaN logits encountered in sample()");
     }
 
+    if !history.is_empty() && (params.presence_penalty != 0.0 || params.repetition_penalty != 1.0) {
+        let seen: HashSet<u32> = history.iter().copied().collect();
+        for id in seen {
+            let i = id as usize;
+            if i >= last_row.len() {
+                continue;
+            }
+            if params.repetition_penalty != 1.0 {
+                let value = last_row[i];
+                last_row[i] = if value > 0.0 {
+                    value / params.repetition_penalty
+                } else {
+                    value * params.repetition_penalty
+                };
+            }
+            last_row[i] -= params.presence_penalty;
+        }
+    }
+
+    if params.temperature == 0.0 {
+        return argmax(&last_row);
+    }
+
     // Apply temperature
     let mut logits_f32: Vec<f32> = last_row.iter().map(|v| v / params.temperature).collect();
 
@@ -195,3 +227,25 @@ fn argmax(data: &[f32]) -> u32 {
     }
     best_i as u32
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn greedy_sampling_applies_history_penalties() {
+        let logits = Tensor::from_slice(&[10.0f32, 9.0, 0.0], &[1, 3]);
+
+        let presence = SamplingParams {
+            presence_penalty: 2.0,
+            ..SamplingParams::default()
+        };
+        assert_eq!(sample_with_history(&logits, &presence, &[0]), 1);
+
+        let repetition = SamplingParams {
+            repetition_penalty: 2.0,
+            ..SamplingParams::default()
+        };
+        assert_eq!(sample_with_history(&logits, &repetition, &[0]), 1);
+    }
+}